"""Parsing and logo-cropping for the general image scan workflow."""

import base64
import io
import re
import uuid
from typing import Optional

from PIL import Image
from pydantic import ValidationError

from src.models.api.image_scan import (
    CATEGORY_POSTER,
    CATEGORY_VISITING_CARD,
    ImageScanLLMOutput,
    ImageScanResponse,
    ScanBBox,
    ScanDocumentMetadata,
    ScanDocumentOut,
    ScannedDocument,
    ScanPoster,
)
from src.models.api.response import ExtractionMetadata
from src.processing.shared.json_utils import parse_json_from_content

LOGO_MIN_SIDE_PX = 16
LOGO_MAX_SIDE_PX = 512
LOGO_PAD_RATIO = 0.1
LOGO_ROUGH_PAD_RATIO = 0.5
LOGO_REFINED_PAD_RATIO = 0.04


def parse_image_scan_response(content: str) -> ImageScanLLMOutput:
    """Parse model output. Raises ValueError when it is not usable JSON."""

    data = parse_json_from_content(content)
    if data is None:
        raise ValueError("Model output was empty or not valid JSON.")
    try:
        return ImageScanLLMOutput.model_validate(data)
    except ValidationError as exc:
        raise ValueError(f"Model output did not match the expected shape: {exc.error_count()} error(s).") from exc


def expand_box(bbox: ScanBBox, pad_ratio: float) -> ScanBBox:
    """Grow a box by a fraction of its own size on every side, kept inside the image."""

    left = max(0.0, bbox.x - bbox.w * pad_ratio)
    top = max(0.0, bbox.y - bbox.h * pad_ratio)
    right = min(1.0, bbox.x + bbox.w + bbox.w * pad_ratio)
    bottom = min(1.0, bbox.y + bbox.h + bbox.h * pad_ratio)
    return ScanBBox(x=left, y=top, w=right - left, h=bottom - top)


def crop_region(image_png: bytes, region: ScanBBox) -> Optional[bytes]:
    """Cut a normalized region out of an image as PNG bytes (no padding, no shrinking)."""

    try:
        image = Image.open(io.BytesIO(image_png))
        image.load()
    except Exception:
        return None
    width, height = image.size
    left, top = int(round(region.x * width)), int(round(region.y * height))
    right, bottom = int(round((region.x + region.w) * width)), int(round((region.y + region.h) * height))
    if right - left < LOGO_MIN_SIDE_PX or bottom - top < LOGO_MIN_SIDE_PX:
        return None
    buffer = io.BytesIO()
    image.crop((left, top, right, bottom)).convert("RGB").save(buffer, format="PNG")
    return buffer.getvalue()


def parse_refined_bbox(content: str) -> Optional[ScanBBox]:
    """Read the tight box from the refinement answer. None when the logo was not found."""

    data = parse_json_from_content(content)
    if not data:
        return None
    raw = data.get("bbox", data)
    if not isinstance(raw, dict):
        return None
    try:
        x, y, w, h = (float(raw[key]) for key in ("x", "y", "w", "h"))
    except (KeyError, TypeError, ValueError):
        return None
    x = min(max(x, 0.0), 1.0)
    y = min(max(y, 0.0), 1.0)
    w = min(max(w, 0.0), 1.0 - x)
    h = min(max(h, 0.0), 1.0 - y)
    if w <= 0 or h <= 0:
        return None
    return ScanBBox(x=x, y=y, w=w, h=h)


def map_box_to_image(region: ScanBBox, local: ScanBBox) -> ScanBBox:
    """Convert a box measured inside a cropped region back to whole-image coordinates."""

    return ScanBBox(
        x=region.x + local.x * region.w,
        y=region.y + local.y * region.h,
        w=local.w * region.w,
        h=local.h * region.h,
    )


def crop_logo(
    image_png: bytes,
    bbox: ScanBBox,
    *,
    pad_ratio: float = LOGO_PAD_RATIO,
) -> Optional[bytes]:
    """Cut the logo box out of an image. Returns PNG bytes, or None when the box is unusable."""

    try:
        image = Image.open(io.BytesIO(image_png))
        image.load()
    except Exception:
        return None

    width, height = image.size
    pad_x = bbox.w * pad_ratio
    pad_y = bbox.h * pad_ratio
    left = max(0, int(round((bbox.x - pad_x) * width)))
    top = max(0, int(round((bbox.y - pad_y) * height)))
    right = min(width, int(round((bbox.x + bbox.w + pad_x) * width)))
    bottom = min(height, int(round((bbox.y + bbox.h + pad_y) * height)))

    if right - left < LOGO_MIN_SIDE_PX or bottom - top < LOGO_MIN_SIDE_PX:
        return None

    cropped = image.crop((left, top, right, bottom)).convert("RGB")
    cropped.thumbnail((LOGO_MAX_SIDE_PX, LOGO_MAX_SIDE_PX), Image.LANCZOS)
    buffer = io.BytesIO()
    cropped.save(buffer, format="PNG")
    return buffer.getvalue()


def _snake_key(label: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", label.lower()).strip("_")


def fill_category_defaults(document: ScannedDocument) -> None:
    """Make sure the object promised for the document's category exists.

    The model normally fills it. If it only filled the generic fields, build the
    category object from those so the client never has to guess where the data is.
    """

    if document.category == CATEGORY_VISITING_CARD and not document.card_data:
        data: dict[str, str] = {}
        for item in document.fields:
            if not item.label or not item.value:
                continue
            key = _snake_key(item.label) or "field"
            candidate, n = key, 2
            while candidate in data:
                candidate, n = f"{key}_{n}", n + 1
            data[candidate] = item.value
        document.card_data = data or None
    elif document.category == CATEGORY_POSTER and document.poster is None:
        document.poster = ScanPoster(subject=document.title or document.summary, details=list(document.fields))


def normalize_image_indexes(output: ImageScanLLMOutput, image_count: int) -> None:
    """Keep only valid image numbers on documents and logos (in place, safe to repeat)."""

    for document in output.documents:
        valid = sorted({i for i in document.image_indexes if 0 <= i < image_count})
        if not valid and image_count == 1:
            valid = [0]
        document.image_indexes = valid
        for logo in document.logos:
            if logo.image_index is None or not 0 <= logo.image_index < image_count:
                logo.image_index = valid[0] if valid else None


def finalize_image_scan(
    output: ImageScanLLMOutput,
    images_png: list[bytes],
    metadata: ExtractionMetadata,
    refined: Optional[set[tuple[int, int]]] = None,
) -> ImageScanResponse:
    """Clean up image indexes and attach cropped logo images.

    `refined` holds (document, logo) positions whose box was tightened by the second pass;
    those get a snug crop, the rest get a roomy one because the first-pass box is rough.
    """

    refined = refined or set()
    image_count = len(images_png)
    normalize_image_indexes(output, image_count)
    for doc_pos, document in enumerate(output.documents):
        fill_category_defaults(document)
        for logo_pos, logo in enumerate(document.logos):
            index = logo.image_index
            if index is None or logo.bbox is None:
                continue
            pad = LOGO_REFINED_PAD_RATIO if (doc_pos, logo_pos) in refined else LOGO_ROUGH_PAD_RATIO
            cropped = crop_logo(images_png[index], logo.bbox, pad_ratio=pad)
            if cropped is not None:
                logo.logo_image_base64 = base64.b64encode(cropped).decode("ascii")
                logo.logo_image_mime = "image/png"

    return ImageScanResponse(
        scan_id=uuid.uuid4().hex,
        image_count=image_count,
        documents=[_to_output(document) for document in output.documents],
        warnings=output.warnings,
        usage=metadata,
    )


def _to_output(document: ScannedDocument) -> ScanDocumentOut:
    return ScanDocumentOut(
        category=document.category,
        document_type=document.document_type,
        issuer=document.issuer,
        card_data=document.card_data,
        poster=document.poster,
        metadata=ScanDocumentMetadata(
            image_indexes=document.image_indexes,
            title=document.title,
            summary=document.summary,
            language=document.language,
            fields=document.fields,
            tables=document.tables,
            full_text=document.full_text,
            logos=document.logos,
        ),
    )