import base64
import io
import json
import time

from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from PIL import Image

from src.config.settings import get_settings

# Maximum dimension (width or height) sent to the model.
# Wide strategic stock sheets need high resolution to read all columns clearly.
_MAX_IMAGE_DIM = 6000


def _parse_json_response(content: str) -> dict:
    text = content.strip()
    if not text:
        raise ValueError("Model returned an empty response")

    # Common case: model wraps JSON in markdown fences.
    if text.startswith("```"):
        lines = text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Fallback: extract the outermost JSON object from mixed text.
    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        candidate = text[start : end + 1]
        return json.loads(candidate)

    raise ValueError(f"Model did not return valid JSON. Raw response: {text[:500]}")


def _resize_image(image_bytes: bytes) -> bytes:
    """Resize image so neither dimension exceeds _MAX_IMAGE_DIM."""
    img = Image.open(io.BytesIO(image_bytes))
    w, h = img.size
    if w <= _MAX_IMAGE_DIM and h <= _MAX_IMAGE_DIM:
        return image_bytes

    scale = _MAX_IMAGE_DIM / max(w, h)
    new_w = int(w * scale)
    new_h = int(h * scale)
    img = img.resize((new_w, new_h), Image.LANCZOS)

    if img.mode in ("RGBA", "P"):
        img = img.convert("RGB")

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def _image_message(image_bytes: bytes, prompt_text: str) -> HumanMessage:
    resized = _resize_image(image_bytes)
    encoded = base64.standard_b64encode(resized).decode("utf-8")
    return HumanMessage(
        content=[
            {"type": "text", "text": prompt_text},
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/png;base64,{encoded}",
                    "detail": "high",
                },
            },
        ]
    )


def invoke_json_extraction(
    system_prompt: str,
    image_bytes: bytes,
    max_tokens_override: int | None = None,
) -> tuple[dict, dict]:
    settings = get_settings()
    # Strategic stock sheets with 43+ rows need much more output space
    effective_max_tokens = max_tokens_override or max(settings.max_tokens, 8192)
    llm = ChatOpenAI(
        model=settings.model_to_use,
        temperature=settings.temperature,
        max_tokens=effective_max_tokens,
        api_key=settings.openai_api_key,
    )

    started_at = time.perf_counter()

    response = llm.invoke(
        [
            SystemMessage(content=system_prompt),
            _image_message(
                image_bytes,
                (
                    "You are looking at a warehouse document image. "
                    "Extract the required fields and return ONLY valid JSON. "
                    "Do not say you cannot process the image — always attempt extraction "
                    "and return null for any field you cannot find."
                ),
            ),
        ]
    )

    latency_ms = round((time.perf_counter() - started_at) * 1000, 2)

    content = response.content if hasattr(response, "content") else str(response)
    if not isinstance(content, str):
        content = str(content)

    # If the model refused to process the image, retry with a stronger prompt
    refusal_phrases = [
        "unable to process",
        "cannot process",
        "can't process",
        "i'm unable",
        "i cannot",
        "please provide the text",
        "provide the details",
    ]
    if any(phrase in content.lower() for phrase in refusal_phrases):
        response = llm.invoke(
            [
                SystemMessage(
                    content=(
                        system_prompt
                        + "\n\nCRITICAL: You MUST attempt to extract fields from the image. "
                        "Even if the image quality is poor, return your best guess for each field. "
                        "Return null only if a field is truly absent. "
                        "NEVER say you cannot process the image."
                    )
                ),
                _image_message(
                    image_bytes,
                    "Extract all fields from this document image. Return only JSON.",
                ),
            ]
        )
        content = response.content if hasattr(response, "content") else str(response)
        if not isinstance(content, str):
            content = str(content)

    parsed = _parse_json_response(content)
    metadata = {
        "model": settings.model_to_use,
        "latency_ms": latency_ms,
    }
    return parsed, metadata