import hashlib
import json
import logging
from pathlib import Path
from typing import Optional, Type

from openai import OpenAI
from pydantic import BaseModel

class NoSchema(BaseModel): ...


logger = logging.getLogger("llm client")

class ResponseCache:
    """File-backed cache: prompt-hash -> validated-model JSON. Makes runs idempotent."""

    def __init__(self, path: Optional[Path] = None) -> None:
        self.path = path
        self._mem: dict[str, str] = {}
        if path and path.exists():
            try:
                self._mem = json.loads(path.read_text("utf-8"))
            except Exception:  # noqa: BLE001
                logger.warning("Corrupt cache at %s; starting fresh", path)

    @staticmethod
    def key(system: str, user: str, schema: Type[BaseModel]) -> str:
        h = hashlib.sha256()
        h.update(schema.__name__.encode())
        h.update(b"\x00")
        h.update(system.encode())
        h.update(b"\x00")
        h.update(user.encode())
        return h.hexdigest()

    def get(self, k: str) -> Optional[str]:
        return self._mem.get(k)

    def put(self, k: str, v: str) -> None:
        self._mem[k] = v
        if self.path:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            tmp = self.path.with_suffix(self.path.suffix + ".tmp")
            tmp.write_text(json.dumps(self._mem, indent=2, sort_keys=True), "utf-8")
            tmp.replace(self.path)
        else:
            logger.warning("No path specified for cache!")


class LLMClient:
    def __init__(self, base_url, api_key: str | None = None, model: str = "gpt-4o-mini", cache: Optional[ResponseCache] = None):
        self.model = model
        self.client = OpenAI(base_url=base_url, api_key=api_key or "ollama")
        self.cache = cache

    def chat_with_schema(self, prompt: str, schema: dict, system: str | None = None) -> str | None:
        # --- idempotency: serve from cache when prompt+system+schema match ---

        cache_key = ResponseCache.key(system or "", prompt, NoSchema) if self.cache else None
        if cache_key and self.cache and (hit := self.cache.get(cache_key)) is not None:
            return json.loads(hit).get('response', "")

        assert_strict_mode_clean(schema)
        kwargs = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": prompt},
            ] if system else [
                {"role": "user", "content": prompt},
            ],
            "response_format": schema,
        }
        resp = self.client.chat.completions.create(**kwargs)

        if cache_key and self.cache:
            self.cache.put(cache_key, json.dumps({"response": resp.choices[0].message.content or ""}))

        return resp.choices[0].message.content

    def chat(self, system: str, user: str, json_mode: bool = False, images: list[str] | None = None) -> str | None:
        images = images or []
        if images:
            content = [{"type": "text", "text": user}]
            for b64 in images:
                content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}"},
                })
        else:
            content = user

        kwargs = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": content},
            ],
        }
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}
        resp = self.client.chat.completions.create(**kwargs)
        return resp.choices[0].message.content


def assert_strict_mode_clean(schema: dict, path: str = "$") -> None:
    """Recursively verify a JSON Schema obeys OpenAI strict mode."""
    if schema.get("type") == "object":
        props = schema.get("properties", {})
        required = set(schema.get("required", []))

        if schema.get("additionalProperties") is not False:
            raise ValueError(f"{path}: additionalProperties must be false")

        for name in props:
            if name not in required:
                raise ValueError(f"{path}.{name}: not in required")

        for name, sub in props.items():
            assert_strict_mode_clean(sub, f"{path}.{name}")

    if schema.get("type") == "array":
        assert_strict_mode_clean(schema["items"], f"{path}[]")

    for name, sub in schema.get("$defs", {}).items():
        assert_strict_mode_clean(sub, f"$defs.{name}")

    # OpenAI rejects these keywords in strict mode
    for banned in ("default", "oneOf"):
        if banned in schema:
            raise ValueError(f"{path}: '{banned}' not allowed in strict mode")