# retoor <retoor@molodetz.nl>
from __future__ import annotations

import asyncio
import base64
import hashlib
import io
import logging
import mimetypes
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

from PIL import Image, UnidentifiedImageError

from devplacepy.services.jobs.isslop.agent.llm import LlmClient, LlmUnavailableError, extract_json_object
from devplacepy.services.jobs.isslop.config import (
    IMAGE_CONCURRENCY,
    IMAGE_EXTENSIONS,
    IMAGE_MAX_BYTES,
    IMAGE_MAX_COUNT,
    IMAGE_MIN_BYTES,
    IMAGE_MIN_DIMENSION,
    THUMBNAIL_MAX_DIMENSION,
    THUMBNAIL_QUALITY,
    VISION_TIMEOUT_SECONDS,
)

logger = logging.getLogger(__name__)

VISION_PROMPT: str = (
    "You are a forensic image analyst deciding whether an image is AI-generated (or heavily AI-edited) or a genuine "
    "human-made photograph, illustration or screenshot. Examine, in detail: hands and fingers, faces and teeth, skin "
    "texture (waxy or plastic over-smoothing), eyes and reflections, text and lettering (garbled or nonsensical), "
    "lighting and shadow direction consistency, background warping and impossible geometry, repeated or cloned patterns, "
    "over-saturated or airbrushed rendering, and the 'too perfect' mathematically smooth look typical of diffusion models. "
    "Screenshots, logos, diagrams, icons and real product photos are usually not AI-generated; judge accordingly and do "
    "not call an ordinary screenshot AI art. Respond with a single JSON object only, using exactly these keys: "
    '{"ai_probability": <integer 0-100, likelihood the image is AI-generated>, '
    '"verdict": <one of "ai-generated"|"likely-ai"|"uncertain"|"likely-human"|"human-made">, '
    '"image_kind": <short phrase, e.g. "photo of a person", "screenshot", "logo", "illustration">, '
    '"tells": [<up to five short strings naming concrete artifacts or reasons>], '
    '"description": <two to four factual sentences describing the image and the evidence>}'
)

VERDICT_VALUES: frozenset[str] = frozenset(
    {"ai-generated", "likely-ai", "uncertain", "likely-human", "human-made"}
)


@dataclass(frozen=True)
class ImageVerdict:
    relative: str
    ai_probability: float
    grade: str
    verdict: str
    image_kind: str
    tells: list[str]
    description: str


def _grade_for(ai_probability: float) -> str:
    if ai_probability <= 15:
        return "A"
    if ai_probability <= 35:
        return "B"
    if ai_probability <= 55:
        return "C"
    if ai_probability <= 75:
        return "D"
    return "F"


def _eligible(path: Path) -> bool:
    if path.suffix.lower() not in IMAGE_EXTENSIONS:
        return False
    try:
        size = path.stat().st_size
    except OSError:
        return False
    if size < IMAGE_MIN_BYTES or size > IMAGE_MAX_BYTES:
        return False
    try:
        with Image.open(path) as image:
            width, height = image.size
    except (UnidentifiedImageError, OSError, ValueError) as error:
        logger.debug("Skipping unreadable image %s: %s", path, error)
        return False
    if width < IMAGE_MIN_DIMENSION or height < IMAGE_MIN_DIMENSION:
        return False
    return True


def collect_images(workspace: Path, limit: int = IMAGE_MAX_COUNT) -> list[Path]:
    candidates = [
        path
        for path in sorted(workspace.rglob("*"))
        if path.is_file() and not path.is_symlink() and _eligible(path)
    ]
    if len(candidates) <= limit:
        logger.info("Collected %d images for vision analysis", len(candidates))
        return candidates
    ranked = sorted(
        candidates,
        key=lambda path: hashlib.sha256(str(path.relative_to(workspace)).encode("utf-8")).hexdigest(),
    )
    sample = sorted(ranked[:limit])
    logger.info("Sampled %d of %d images deterministically for vision analysis", len(sample), len(candidates))
    return sample


def _to_data_url(path: Path) -> Optional[str]:
    mime_type = mimetypes.guess_type(path.name)[0]
    if not mime_type or not mime_type.startswith("image/"):
        return None
    try:
        raw = path.read_bytes()
    except OSError as error:
        logger.warning("Image read failed for %s: %s", path, error)
        return None
    return f"data:{mime_type};base64,{base64.b64encode(raw).decode('utf-8')}"


async def classify_image(llm: LlmClient, workspace: Path, path: Path) -> Optional[ImageVerdict]:
    relative = str(path.relative_to(workspace))
    data_url = await asyncio.to_thread(_to_data_url, path)
    if data_url is None:
        return None
    try:
        answer = await llm.describe_image(VISION_PROMPT, data_url, VISION_TIMEOUT_SECONDS)
    except LlmUnavailableError as error:
        logger.warning("Vision unavailable for %s: %s", relative, error)
        return None
    parsed = extract_json_object(answer)
    if parsed is None:
        logger.warning("Vision returned unparseable output for %s", relative)
        return None
    probability = parsed.get("ai_probability")
    if not isinstance(probability, (int, float)) or isinstance(probability, bool):
        probability = 50.0
    probability = float(min(100.0, max(0.0, probability)))
    verdict = str(parsed.get("verdict", "uncertain")).strip().lower()
    if verdict not in VERDICT_VALUES:
        verdict = "uncertain"
    tells_field = parsed.get("tells")
    tells = [str(item)[:120] for item in tells_field[:5]] if isinstance(tells_field, list) else []
    return ImageVerdict(
        relative=relative,
        ai_probability=probability,
        grade=_grade_for(probability),
        verdict=verdict,
        image_kind=str(parsed.get("image_kind", "image"))[:80],
        tells=tells,
        description=str(parsed.get("description", ""))[:1000],
    )


async def analyze_images(
    llm: LlmClient,
    workspace: Path,
    images: list[Path],
) -> list[ImageVerdict]:
    semaphore = asyncio.Semaphore(IMAGE_CONCURRENCY)

    async def worker(path: Path) -> Optional[ImageVerdict]:
        async with semaphore:
            return await classify_image(llm, workspace, path)

    results = await asyncio.gather(*(worker(path) for path in images))
    verdicts = [result for result in results if result is not None]
    logger.info("Vision analysis produced %d verdicts from %d images", len(verdicts), len(images))
    return verdicts


def make_thumbnail(source: Path, media_dir: Path, relative: str) -> Optional[str]:
    name = f"{hashlib.sha1(relative.encode('utf-8')).hexdigest()[:16]}.webp"
    try:
        media_dir.mkdir(parents=True, exist_ok=True)
        with Image.open(source) as image:
            image.thumbnail((THUMBNAIL_MAX_DIMENSION, THUMBNAIL_MAX_DIMENSION))
            if image.mode not in ("RGB", "RGBA"):
                image = image.convert("RGBA")
            image.save(media_dir / name, format="WEBP", quality=THUMBNAIL_QUALITY)
    except (UnidentifiedImageError, OSError, ValueError) as error:
        logger.warning("Thumbnail failed for %s: %s", relative, error)
        return None
    return name


def persist_screenshot(raw_bytes: bytes, media_dir: Path, relative: str) -> Optional[str]:
    name = f"{hashlib.sha1(relative.encode('utf-8')).hexdigest()[:16]}.webp"
    try:
        media_dir.mkdir(parents=True, exist_ok=True)
        with Image.open(io.BytesIO(raw_bytes)) as image:
            image.thumbnail((THUMBNAIL_MAX_DIMENSION, THUMBNAIL_MAX_DIMENSION))
            if image.mode not in ("RGB", "RGBA"):
                image = image.convert("RGBA")
            image.save(media_dir / name, format="WEBP", quality=THUMBNAIL_QUALITY)
    except (UnidentifiedImageError, OSError, ValueError) as error:
        logger.warning("Screenshot persist failed for %s: %s", relative, error)
        return None
    return name


def image_summary(verdicts: list[ImageVerdict]) -> dict[str, float | int | str]:
    if not verdicts:
        return {"count": 0, "mean_ai_probability": 0.0, "grade": "n/a", "ai_generated_count": 0}
    mean = sum(verdict.ai_probability for verdict in verdicts) / len(verdicts)
    ai_generated = sum(1 for verdict in verdicts if verdict.ai_probability >= 60.0)
    return {
        "count": len(verdicts),
        "mean_ai_probability": round(mean, 1),
        "grade": _grade_for(mean),
        "ai_generated_count": ai_generated,
    }
