# retoor <retoor@molodetz.nl>

import asyncio
import hashlib
import logging
import time
from collections import OrderedDict
from typing import Any, Optional

import httpx

from devplacepy.services.openai_gateway.config import VISION_INSTRUCTION
from devplacepy.services.openai_gateway.thinking import apply_thinking
from devplacepy.services.openai_gateway.usage import classify_error

logger = logging.getLogger(__name__)


class VisionCache:
    def __init__(self, size: int = 256):
        self.size = size
        self._store: "OrderedDict[str, str]" = OrderedDict()

    def get(self, key: str) -> Optional[str]:
        if self.size <= 0:
            return None
        value = self._store.get(key)
        if value is not None:
            self._store.move_to_end(key)
        return value

    def put(self, key: str, value: str) -> None:
        if self.size <= 0:
            return
        self._store[key] = value
        self._store.move_to_end(key)
        while len(self._store) > self.size:
            self._store.popitem(last=False)

    def __len__(self) -> int:
        return len(self._store)


def has_vision_blocks(content: Any) -> bool:
    if not isinstance(content, list):
        return False
    return any(
        isinstance(b, dict) and b.get("type") in ("image_url", "image") for b in content
    )


def split_text_and_images(content: list) -> tuple[str, list]:
    texts: list[str] = []
    images: list[dict] = []
    for block in content:
        if not isinstance(block, dict):
            continue
        btype = block.get("type")
        if btype == "text":
            t = block.get("text", "")
            if t:
                texts.append(t)
        elif btype in ("image_url", "image"):
            images.append(block)
    return "\n".join(texts).strip(), images


def _image_cache_key(image_block: dict) -> str:
    iu = image_block.get("image_url") or image_block.get("image") or ""
    url = iu.get("url", "") if isinstance(iu, dict) else str(iu)
    return hashlib.sha256(url.encode("utf-8", errors="replace")).hexdigest()[:32]


def _format_vision_block(descriptions: list) -> str:
    if len(descriptions) == 1:
        return f"[Image seen by vision model:\n{descriptions[0]}\n]"
    parts = [f"Image {i}:\n{d}" for i, d in enumerate(descriptions, 1)]
    return (
        f"[{len(descriptions)} images seen by vision model:\n"
        + "\n\n".join(parts)
        + "\n]"
    )


class VisionAugmenter:
    def __init__(
        self,
        vision_url: str,
        vision_model: str,
        vision_key: str,
        cache: VisionCache,
        referer: str = "",
        title: str = "",
        vision_dialect: str = "",
        ledger=None,
        owner: tuple = ("unknown", "unknown"),
        pricing=None,
        context_map=None,
        app_reference: str = "default",
    ):
        self.vision_url = vision_url
        self.vision_model = vision_model
        self.vision_key = vision_key
        self.cache = cache
        self.referer = referer
        self.title = title
        self.vision_dialect = vision_dialect
        self.ledger = ledger
        self.owner = owner
        self.pricing = pricing
        self.context_map = context_map or {}
        self.app_reference = app_reference
        self.calls = 0
        self.cost_usd = 0.0

    def _record(self, latency_ms, status_code, success, category, usage):
        if self.ledger is None or self.pricing is None:
            return
        row = self.ledger.record(
            {
                "owner_kind": self.owner[0],
                "owner_id": self.owner[1],
                "backend": "vision",
                "endpoint": "chat/completions",
                "model": self.vision_model,
                "requested_model": self.vision_model,
                "temperature": 0.2,
                "upstream_latency_ms": latency_ms,
                "total_latency_ms": latency_ms,
                "status_code": status_code,
                "success": success,
                "error_category": category,
                "usage": usage,
                "app_reference": self.app_reference,
            },
            self.pricing,
            self.context_map,
        )
        if row:
            self.cost_usd += float(row.get("cost_usd") or 0.0)

    async def _describe_one(self, client: httpx.AsyncClient, image_block: dict) -> str:
        if not self.vision_key:
            return "[vision unavailable: vision API key not configured]"
        payload = {
            "model": self.vision_model,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": VISION_INSTRUCTION},
                        image_block,
                    ],
                }
            ],
            "temperature": 0.2,
            "stream": False,
        }
        apply_thinking(payload, self.vision_url, default_enabled=False, dialect=self.vision_dialect)
        headers = {
            "Authorization": f"Bearer {self.vision_key}",
            "Content-Type": "application/json",
        }
        if self.referer:
            headers["HTTP-Referer"] = self.referer
        if self.title:
            headers["X-Title"] = self.title
        start = time.monotonic()
        try:
            resp = await client.post(self.vision_url, json=payload, headers=headers)
        except httpx.RequestError as e:
            logger.warning("vision connection failed: %s", e)
            self._record(
                (time.monotonic() - start) * 1000,
                502,
                False,
                classify_error(0, e),
                None,
            )
            return f"[vision call failed: {e}]"
        latency_ms = (time.monotonic() - start) * 1000
        if resp.status_code != 200:
            logger.warning("vision %s: %s", resp.status_code, resp.text[:200])
            self._record(
                latency_ms,
                resp.status_code,
                False,
                classify_error(resp.status_code, None, resp.text),
                None,
            )
            return f"[vision failed: HTTP {resp.status_code}]"
        try:
            data = resp.json()
            text = data["choices"][0]["message"].get("content") or ""
            self._record(latency_ms, 200, True, None, data.get("usage"))
            return text.strip() or "[vision returned empty response]"
        except (KeyError, IndexError, ValueError) as e:
            self._record(latency_ms, 200, False, "gateway", None)
            return f"[vision parse error: {e}]"

    async def _describe_images(self, client: httpx.AsyncClient, images: list) -> list:
        keys = [_image_cache_key(img) for img in images]
        descriptions: list = [self.cache.get(k) for k in keys]
        miss = [i for i, d in enumerate(descriptions) if d is None]
        if miss:
            self.calls += len(miss)
            results = await asyncio.gather(
                *[self._describe_one(client, images[i]) for i in miss]
            )
            for idx, result in zip(miss, results):
                descriptions[idx] = result
                if not result.startswith("[vision"):
                    self.cache.put(keys[idx], result)
        return [d or "" for d in descriptions]

    async def augment_messages(self, client: httpx.AsyncClient, messages: list) -> list:
        if not any(
            has_vision_blocks(m.get("content")) for m in messages if isinstance(m, dict)
        ):
            return messages
        out: list = []
        for m in messages:
            content = m.get("content") if isinstance(m, dict) else None
            if not has_vision_blocks(content):
                out.append(m)
                continue
            user_text, images = split_text_and_images(content)
            descriptions = await self._describe_images(client, images)
            vision_block = _format_vision_block(descriptions)
            merged = f"{user_text}\n\n{vision_block}" if user_text else vision_block
            new_msg = dict(m)
            new_msg["content"] = merged
            out.append(new_msg)
        return out
