"""Vision extraction via Claude.

Two calls:
  - extract_lead(): screenshots + a recipe's schema -> raw structured JSON.
  - classify_page(): one screenshot -> {site, page_type, address bar url}, used
    only when no native browser URL is available (visual site detection).

Both use structured outputs (output_config.format) so the response is
schema-validated JSON, not free text.
"""
from __future__ import annotations

import base64
import json
from functools import lru_cache
from typing import Any

from config import ANTHROPIC_MODEL, MAX_IMAGES_PER_CALL
from ai.recipes.base import SiteRecipe


@lru_cache(maxsize=1)
def _client():
    import anthropic

    # Resolves ANTHROPIC_API_KEY (or an `ant auth login` profile) from the env.
    return anthropic.Anthropic()


def _image_block(path: str) -> dict[str, Any]:
    with open(path, "rb") as f:
        data = base64.standard_b64encode(f.read()).decode("utf-8")
    return {
        "type": "image",
        "source": {"type": "base64", "media_type": "image/png", "data": data},
    }


def _first_json(response) -> dict[str, Any]:
    for block in response.content:
        if block.type == "text":
            return json.loads(block.text)
    raise RuntimeError("model returned no text block")


def extract_lead(recipe: SiteRecipe, screenshot_paths: list[str]) -> dict[str, Any]:
    paths = screenshot_paths[:MAX_IMAGES_PER_CALL]
    content: list[dict[str, Any]] = [_image_block(p) for p in paths]
    content.append({"type": "text", "text": recipe.extraction_prompt})

    response = _client().messages.create(
        model=ANTHROPIC_MODEL,
        max_tokens=2048,
        messages=[{"role": "user", "content": content}],
        output_config={"format": {"type": "json_schema", "schema": recipe.extraction_schema}},
    )
    return _first_json(response)


_CLASSIFY_SCHEMA = {
    "type": "object",
    "properties": {
        "site": {"type": "string"},        # e.g. "linkedin", "other"
        "page_type": {"type": "string"},   # e.g. "profile", "search", "feed", "other"
        "address_bar_url": {"type": "string"},
    },
    "required": ["site", "page_type", "address_bar_url"],
    "additionalProperties": False,
}

_CLASSIFY_PROMPT = (
    "This is a screenshot of a web browser. Identify the site and page from the "
    "browser's address bar and the visible layout. Return: 'site' (lowercase, "
    "e.g. 'linkedin', or 'other'); 'page_type' ('profile' for a single person's "
    "page, otherwise 'search', 'feed', 'messaging', or 'other'); and "
    "'address_bar_url' (the URL text you can read, or \"\")."
)


def classify_page(screenshot_path: str) -> dict[str, Any]:
    response = _client().messages.create(
        model=ANTHROPIC_MODEL,
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": [_image_block(screenshot_path), {"type": "text", "text": _CLASSIFY_PROMPT}],
        }],
        output_config={"format": {"type": "json_schema", "schema": _CLASSIFY_SCHEMA}},
    )
    return _first_json(response)