#!/usr/bin/env python3
"""RouterAI Image MCP server.

Даёт ЛЮБОЙ модели (в т.ч. той, что не умеет рисовать) инструмент
`generate_image` — генерация/редактирование картинок через RouterAI
(Nano Banana / Gemini Image, FLUX, GPT-Image).

Транспорт:
  MCP_TRANSPORT=stdio (по умолчанию) | http
  MCP_PORT=8790 (для http)

Ключ ищем в порядке:
  1) env ROUTERAI_API_KEY / IMAGE_MCP_API_KEY
  2) ~/.hermes/models.yaml -> models.deepseek.api_key (base_url = routerai.ru)
  3) ~/.hermes/models.yaml -> любой api_key с base_url routerai.ru
"""

from __future__ import annotations

import base64
import json
import os
import re
import sys
import time
import mimetypes
import urllib.error
import urllib.request
from pathlib import Path

from mcp.server.fastmcp import FastMCP

BASE_URL = os.environ.get("ROUTERAI_BASE_URL", "https://routerai.ru/api/v1").rstrip("/")
DEFAULT_MODEL = os.environ.get("IMAGE_MCP_DEFAULT_MODEL", "google/gemini-3.1-flash-image")
DEFAULT_OUT_DIR = os.environ.get(
    "IMAGE_MCP_OUT_DIR", str(Path.home() / ".hermes" / "image_cache" / "mcp-gen")
)
REQUEST_TIMEOUT = int(os.environ.get("IMAGE_MCP_TIMEOUT", "300"))

# модель -> (бэкенд, человеческое имя). backend "chat" = OpenAI chat с modalities,
# backend "images" = POST /images/generations
MODELS: dict[str, tuple[str, str]] = {
    "google/gemini-3.1-flash-image": ("chat", "Nano Banana 2 (Gemini 3.1 Flash Image)"),
    "google/gemini-3.1-flash-lite-image": ("chat", "Nano Banana 2 Lite"),
    "google/gemini-2.5-flash-image": ("chat", "Nano Banana (Gemini 2.5 Flash Image)"),
    "google/gemini-3-pro-image": ("chat", "Nano Banana Pro (Gemini 3 Pro Image)"),
    "google/gemini-3.1-flash-image-preview": ("chat", "Nano Banana 2 (preview)"),
    "google/gemini-3-pro-image-preview": ("chat", "Nano Banana Pro (preview)"),
    "openai/gpt-image-1": ("images", "GPT Image 1"),
    "openai/gpt-image-1-mini": ("images", "GPT Image 1 Mini"),
    "black-forest-labs/flux.2-pro": ("images", "FLUX.2 Pro"),
    "black-forest-labs/flux.2-max": ("images", "FLUX.2 Max"),
    "black-forest-labs/flux.2-flex": ("images", "FLUX.2 Flex"),
    "black-forest-labs/flux.2-klein-4b": ("images", "FLUX.2 Klein 4B"),
}

mcp = FastMCP("routerai-image")


# --------------------------------------------------------------------------- key


def _api_key() -> str:
    for env in ("ROUTERAI_API_KEY", "IMAGE_MCP_API_KEY"):
        v = os.environ.get(env)
        if v and len(v) > 20:
            return v.strip()

    models_yaml = Path.home() / ".hermes" / "models.yaml"
    if models_yaml.exists():
        try:
            import yaml  # type: ignore

            data = yaml.safe_load(models_yaml.read_text(encoding="utf-8")) or {}
            cands = []
            for name, entry in (data.get("models") or {}).items():
                if not isinstance(entry, dict):
                    continue
                key = str(entry.get("api_key") or "")
                base = str(entry.get("base_url") or "")
                if "routerai.ru" in base and key.startswith("sk-"):
                    cands.append((name != "deepseek", key))
            if cands:
                cands.sort()
                return cands[0][1]
        except Exception as exc:  # noqa: BLE001
            print(f"[routerai-image] models.yaml parse failed: {exc}", file=sys.stderr)

    raise RuntimeError(
        "RouterAI API key not found. Set ROUTERAI_API_KEY, or put a routerai.ru "
        "entry with api_key into ~/.hermes/models.yaml"
    )


# ---------------------------------------------------------------------------- http


def _post(path: str, body: dict, timeout: int = REQUEST_TIMEOUT) -> dict:
    key = _api_key()
    data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    last_err = ""

    for attempt in range(4):
        req = urllib.request.Request(
            f"{BASE_URL}{path}",
            data=data,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {key}",
                "User-Agent": "routerai-image-mcp/1.0",
            },
        )
        try:
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                raw = resp.read()
            # RouterAI иногда отдаёт мусорные префиксы перед JSON
            start = raw.find(b"{")
            return json.loads(raw[start:] if start > 0 else raw)
        except urllib.error.HTTPError as exc:
            payload = exc.read().decode("utf-8", "replace")[:400]
            last_err = f"HTTP {exc.code}: {payload}"
            if exc.code in (429, 500, 502, 503, 504) and attempt < 3:
                time.sleep(3 * (attempt + 1))
                continue
            raise RuntimeError(last_err) from None
        except Exception as exc:  # noqa: BLE001
            last_err = f"{type(exc).__name__}: {exc}"
            if attempt < 3:
                time.sleep(3 * (attempt + 1))
                continue
            raise RuntimeError(last_err) from None
    raise RuntimeError(last_err)


# ------------------------------------------------------------------------- helpers


def _data_url(path: str) -> str:
    p = Path(path).expanduser()
    if not p.exists():
        raise RuntimeError(f"reference image not found: {p}")
    mime = mimetypes.guess_type(p.name)[0] or "image/png"
    return f"data:{mime};base64,{base64.b64encode(p.read_bytes()).decode()}"


def _decode_image(item: dict) -> bytes:
    """Принимает {'b64_json': ...} или {'url': 'data:...'} / http url."""
    if item.get("b64_json"):
        return base64.b64decode(item["b64_json"])
    url = item.get("url") or (item.get("image_url") or {}).get("url") or ""
    if url.startswith("data:"):
        return base64.b64decode(url.split(",", 1)[1])
    if url.startswith("http"):
        with urllib.request.urlopen(url, timeout=REQUEST_TIMEOUT) as resp:  # noqa: S310
            return resp.read()
    raise RuntimeError("no image payload in response")


def _ext_from_bytes(blob: bytes) -> str:
    if blob[:8] == b"\x89PNG\r\n\x1a\n":
        return ".png"
    if blob[:3] == b"\xff\xd8\xff":
        return ".jpg"
    if blob[:4] == b"RIFF" and blob[8:12] == b"WEBP":
        return ".webp"
    return ".png"


def _slug(text: str, limit: int = 40) -> str:
    txt = re.sub(r"[^\w\s-]", "", text, flags=re.UNICODE).strip()
    txt = re.sub(r"[\s-]+", "-", txt)
    return (txt[:limit] or "image").strip("-").lower()


def _run(model: str, prompt: str, refs: list[str] | None, out_dir: str, filename: str | None) -> str:
    if model not in MODELS:
        known = ", ".join(sorted(MODELS))
        raise RuntimeError(f"unknown model '{model}'. Known: {known}")
    backend, human = MODELS[model]

    refs = refs or []
    started = time.time()

    if backend == "chat":
        if refs:
            content: list[dict] = [{"type": "text", "text": prompt}]
            content += [
                {"type": "image_url", "image_url": {"url": _data_url(p)}} for p in refs
            ]
        else:
            content = prompt  # type: ignore[assignment]
        body = {
            "model": model,
            "messages": [{"role": "user", "content": content}],
            "modalities": ["image", "text"],
        }
        if not refs:
            body["response_modalities"] = ["IMAGE"]
        data = _post("/chat/completions", body)
        msg = (data.get("choices") or [{}])[0].get("message", {}) or {}
        images = msg.get("images") or []
        if not images and isinstance(msg.get("content"), list):
            images = [c for c in msg["content"] if c.get("type") in ("image_url", "image")]
        if not images:
            raise RuntimeError(f"model returned no image. raw: {json.dumps(data)[:400]}")
        blob = _decode_image(images[0])
        usage = data.get("usage") or {}
    else:
        body = {"model": model, "prompt": prompt, "n": 1}
        if refs:
            body["image"] = _data_url(refs[0])
        data = _post("/images/generations", body)
        items = data.get("data") or []
        if not items:
            raise RuntimeError(f"model returned no image. raw: {json.dumps(data)[:400]}")
        blob = _decode_image(items[0])
        usage = data.get("usage") or {}

    out = Path(out_dir).expanduser()
    out.mkdir(parents=True, exist_ok=True)
    name = filename or f"{time.strftime('%Y%m%d-%H%M%S')}-{_slug(prompt)}"
    if not Path(name).suffix:
        name += _ext_from_bytes(blob)
    dest = out / name
    dest.write_bytes(blob)

    cost = usage.get("cost")
    secs = time.time() - started
    cost_line = f"{cost:.2f}₽" if isinstance(cost, (int, float)) else "n/a"
    kb = len(blob) / 1024
    return (
        f"Готово: {human}\n"
        f"Файл: {dest} ({kb:.0f} KB)\n"
        f"Модель: {model} | время: {secs:.1f}s | стоимость: {cost_line}\n"
        f"MEDIA:{dest}\n"
        f"(покажи файл пользователю: вставь строку MEDIA:<путь> целиком в ответ)"
    )


# --------------------------------------------------------------------------- tools


@mcp.tool()
def generate_image(
    prompt: str,
    model: str = DEFAULT_MODEL,
    reference_images: list[str] | None = None,
    out_dir: str = DEFAULT_OUT_DIR,
    filename: str | None = None,
) -> str:
    """Сгенерировать картинку по текстовому ТЗ (или отредактировать существующую).

    Args:
        prompt: ТЗ на картинке — что нарисовать. Чем конкретнее (сюжет, стиль,
            свет, композиция, пропорции), тем лучше результат.
        model: id модели RouterAI. Дефолт google/gemini-3.1-flash-image
            (Nano Banana 2). Дешёвые: google/gemini-3.1-flash-lite-image,
            google/gemini-2.5-flash-image. Дорогие: google/gemini-3-pro-image.
            Альтернативы: black-forest-labs/flux.2-pro, openai/gpt-image-1.
        reference_images: список путей к картинкам-референсам/для правки.
        out_dir: куда сложить результат (по умолчанию ~/.hermes/image_cache/mcp-gen).
        filename: имя файла без пути (расширение подставится само).

    Returns:
        Текст с абсолютным путём к файлу и строкой `MEDIA:<путь>` — вставь эту
        строку в свой ответ, чтобы клиент показал картинку.
    """
    return _run(model, prompt, reference_images, out_dir, filename)


@mcp.tool()
def list_image_models() -> str:
    """Список доступных моделей генерации картинок и их бэкендов."""
    lines = [f"{mid} — {human} (backend: {backend})" for mid, (backend, human) in sorted(MODELS.items())]
    return "Доступные модели:\n" + "\n".join(lines) + f"\n\nДефолт: {DEFAULT_MODEL}\nБазовый URL: {BASE_URL}"


def main() -> None:
    transport = os.environ.get("MCP_TRANSPORT", "stdio").lower()
    if transport in ("http", "streamable-http", "sse"):
        mcp.settings.host = os.environ.get("MCP_HOST", "0.0.0.0")
        mcp.settings.port = int(os.environ.get("MCP_PORT", "8790"))
        mcp.run(transport="streamable-http")
    else:
        mcp.run(transport="stdio")


if __name__ == "__main__":
    main()