"""Step 5: Gemini annotation of checkbox locations and states."""

import base64
import io
import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, List, Literal

from pydantic import BaseModel, Field

from .config import (
    ANNOTATION_PROMPT_TEMPLATE,
    GEMINI_MODEL,
    GEMINI_RPM_LIMIT,
    IMAGES_DIR,
    MAX_IMAGE_LONGEST_SIDE,
    PAGE_ANNOTS_DIR,
)
from .download import get_pdf_path
from .state import PipelineState

logger = logging.getLogger(__name__)


# ── Pydantic schema for structured output ────────────────────────────────────

class CheckboxAnnotation(BaseModel):
    """A single checkbox detection."""
    label: Literal["checkbox_checked", "checkbox_unchecked"]
    box_2d: List[int] = Field(
        description="Bounding box as [ymin, xmin, ymax, xmax] normalized to 0-1000"
    )


class CheckboxAnnotations(BaseModel):
    """All checkbox detections on a page."""
    checkboxes: List[CheckboxAnnotation]


# ── Annotation logic ─────────────────────────────────────────────────────────

def _get_openai_client():
    """Get an OpenAI-compatible client for Gemini."""
    try:
        from openai import OpenAI
    except ImportError:
        raise ImportError("openai package required. Install with: pip install openai")

    import os
    from dotenv import load_dotenv
    load_dotenv()

    api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError(
            "Set GEMINI_API_KEY or GOOGLE_API_KEY environment variable"
        )

    return OpenAI(
        api_key=api_key,
        base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
    )


def _annot_path(doc_id: str, page_idx: int) -> Path:
    """Path for a per-page annotation JSON file."""
    return PAGE_ANNOTS_DIR / f"{doc_id}_p{page_idx:03d}.json"


def _image_path(doc_id: str, page_idx: int, split: str = "staging") -> Path:
    """Path for a rendered page PNG."""
    return IMAGES_DIR / split / f"{doc_id}_p{page_idx:03d}.png"


def annotate_page(
    client,
    image_bytes: bytes,
    width: int,
    height: int,
) -> Dict[str, Any]:
    """Send a single page image to Gemini and parse the response.

    Uses structured outputs (Pydantic schema) to guarantee valid JSON from Gemini.

    The image should already be resized so its longest side = MAX_IMAGE_LONGEST_SIDE.
    Gemini returns normalized 0-1000 coords which we convert to pixels here.

    Returns a dict with "checkboxes" containing pixel-space bounding boxes.
    """
    prompt = ANNOTATION_PROMPT_TEMPLATE
    b64_image = base64.b64encode(image_bytes).decode("utf-8")

    response = client.beta.chat.completions.parse(
        model=GEMINI_MODEL,
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/png;base64,{b64_image}",
                        },
                    },
                ],
            }
        ],
        response_format=CheckboxAnnotations,
        max_tokens=16384,
    )
    parsed = response.choices[0].message.parsed
    raw = parsed.model_dump()
    return _normalized_to_pixels(raw, width, height)


def _normalized_to_pixels(raw: Dict[str, Any], width: int, height: int) -> Dict[str, Any]:
    """Convert Gemini's normalized 0-1000 box_2d coords to pixel-space bbox."""
    converted = []
    for cb in raw.get("checkboxes", []):
        box = cb.get("box_2d", [])
        if len(box) != 4:
            continue

        # Gemini format: [ymin, xmin, ymax, xmax] in 0-1000
        ay, ax, by, bx = box
        xmin_n, xmax_n = min(ax, bx), max(ax, bx)
        ymin_n, ymax_n = min(ay, by), max(ay, by)

        pixel_bbox = [
            int(xmin_n * width / 1000),
            int(ymin_n * height / 1000),
            int(xmax_n * width / 1000),
            int(ymax_n * height / 1000),
        ]

        label = cb.get("label", "checkbox_unchecked")
        converted.append({
            "bbox": pixel_bbox,
            "state": "checked" if label == "checkbox_checked" else "unchecked",
        })

    return {"checkboxes": converted}


def annotate(state: PipelineState) -> int:
    """Annotate all qualifying pages with Gemini.

    Returns number of newly annotated pages.
    """
    from tqdm import tqdm

    from natural_pdf import PDF

    PAGE_ANNOTS_DIR.mkdir(parents=True, exist_ok=True)
    client = _get_openai_client()

    # Rate limiting via elapsed-time tracking (not blind sleep)
    min_interval = 60.0 / GEMINI_RPM_LIMIT
    last_request_time = 0.0

    # First pass: collect all pages needing annotation
    work_items: List[tuple] = []  # (doc_id, page_idx)
    for doc_id in state.all_doc_ids():
        if not state.is_done(doc_id, "structural_filtered"):
            continue

        qualifying_pages = state.get(doc_id, "qualifying_pages", [])
        if not qualifying_pages:
            continue

        pdf_path = get_pdf_path(doc_id)
        if not pdf_path.exists():
            continue

        annotated_pages = state.get(doc_id, "annotated_pages", [])
        for page_idx in qualifying_pages:
            if page_idx in annotated_pages:
                continue
            annot_file = _annot_path(doc_id, page_idx)
            if annot_file.exists():
                if page_idx not in annotated_pages:
                    annotated_pages.append(page_idx)
                continue
            work_items.append((doc_id, page_idx))

        # Persist any state fixes from existing-file detection
        state.set(doc_id, "annotated_pages", sorted(set(annotated_pages)))
        if set(qualifying_pages) <= set(annotated_pages):
            state.mark_done(doc_id, "annotated")

    if not work_items:
        logger.info("Annotation: nothing to do (all pages already annotated)")
        return 0

    annotated = 0
    failed = 0
    checked_total = 0
    unchecked_total = 0
    current_pdf = None
    current_doc_id = None

    pbar = tqdm(work_items, desc="Annotating", unit="page")
    for doc_id, page_idx in pbar:
        # Open PDF if needed (cache across pages of same doc)
        if doc_id != current_doc_id:
            if current_pdf is not None:
                current_pdf.close()
            pdf_path = get_pdf_path(doc_id)
            try:
                current_pdf = PDF(str(pdf_path))
                current_doc_id = doc_id
            except Exception as e:
                logger.warning("Failed to open PDF %s: %s", doc_id, e)
                current_pdf = None
                current_doc_id = None
                failed += 1
                continue

        if page_idx >= len(current_pdf.pages):
            continue

        page = current_pdf.pages[page_idx]

        # Render page to PNG at 150 DPI, then resize
        try:
            pil_image = page.to_image(resolution=150)
        except Exception:
            try:
                pil_image = page.render(resolution=150)
            except Exception as e:
                logger.warning("Failed to render doc %s page %d: %s", doc_id, page_idx, e)
                failed += 1
                continue

        # Resize so longest side = MAX_IMAGE_LONGEST_SIDE
        orig_w, orig_h = pil_image.size
        longest = max(orig_w, orig_h)
        if longest > MAX_IMAGE_LONGEST_SIDE:
            scale = MAX_IMAGE_LONGEST_SIDE / longest
            pil_image = pil_image.resize((int(orig_w * scale), int(orig_h * scale)))

        w, h = pil_image.size

        # Save resized PNG to staging
        img_path = _image_path(doc_id, page_idx, split="staging")
        img_path.parent.mkdir(parents=True, exist_ok=True)
        pil_image.save(img_path, format="PNG")

        # Get image bytes for API
        buf = io.BytesIO()
        pil_image.save(buf, format="PNG")
        image_bytes = buf.getvalue()

        # Rate-limit based on elapsed time since last request
        elapsed = time.monotonic() - last_request_time
        if elapsed < min_interval:
            time.sleep(min_interval - elapsed)

        # Annotate with exponential backoff on transient failures
        max_retries = 3
        annotation = None
        for attempt in range(max_retries):
            try:
                last_request_time = time.monotonic()
                annotation = annotate_page(client, image_bytes, w, h)
                break
            except Exception as e:
                if attempt < max_retries - 1:
                    wait = 2 ** (attempt + 1)
                    pbar.write(
                        f"  Attempt {attempt + 1} failed for {doc_id} p{page_idx}: {e} (retry in {wait}s)"
                    )
                    time.sleep(wait)
                else:
                    pbar.write(
                        f"  FAILED {doc_id} p{page_idx} after {max_retries} attempts: {e}"
                    )
                    failed += 1

        if annotation is None:
            continue

        # Save annotation
        annot_file = _annot_path(doc_id, page_idx)
        annot_file.parent.mkdir(parents=True, exist_ok=True)
        with open(annot_file, "w") as f:
            json.dump(annotation, f, indent=2)

        # Track stats
        cbs = annotation.get("checkboxes", [])
        n_checked = sum(1 for c in cbs if c.get("state") == "checked")
        n_unchecked = len(cbs) - n_checked
        checked_total += n_checked
        unchecked_total += n_unchecked
        annotated += 1

        # Update state for this doc
        annotated_pages = state.get(doc_id, "annotated_pages", [])
        annotated_pages.append(page_idx)
        state.set(doc_id, "annotated_pages", sorted(set(annotated_pages)))
        qualifying_pages = state.get(doc_id, "qualifying_pages", [])
        if set(qualifying_pages) <= set(annotated_pages):
            state.mark_done(doc_id, "annotated")

        pbar.set_postfix(ok=annotated, fail=failed, chk=checked_total, unchk=unchecked_total)

    if current_pdf is not None:
        current_pdf.close()

    pbar.close()
    logger.info(
        "Annotation complete. %d new, %d failed. Totals: %d checked, %d unchecked.",
        annotated, failed, checked_total, unchecked_total,
    )
    return annotated


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    annotate(st)
