"""Hermes image backend that keeps owner OAuth outside family chat pods."""

from __future__ import annotations

import base64
import json
import os
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

from agent.image_gen_provider import (
    DEFAULT_ASPECT_RATIO,
    ImageGenProvider,
    error_response,
    normalize_reference_images,
    resolve_aspect_ratio,
    save_b64_image,
    success_response,
)


BROKER_URL = os.environ.get(
    "HERMES_IMAGE_BROKER_URL",
    "http://hermes-image-broker.hermes.svc.cluster.local:9002",
).rstrip("/")
DEFAULT_MODEL = "atlas-image-auto-high"
ROUTE_MODELS = {
    "auto": "atlas-image-auto-high",
    "hosted": "gpt-image-2-high",
    "local": "flux-2-klein-4b-local",
}
MODELS = {
    "atlas-image-auto-high": (
        "Auto (GPT Image quality, local fallback)",
        "auto-high",
        "Hosted quality first; FLUX fallback when hosted generation fails",
    ),
    "gpt-image-2-low": ("GPT Image 2 (Fast)", "low"),
    "gpt-image-2-medium": ("GPT Image 2 (Balanced)", "medium"),
    "gpt-image-2-high": ("GPT Image 2 (Highest quality)", "high"),
    "flux-2-klein-4b-local": (
        "Local FLUX 2 Klein 4B",
        "local-high",
        "Private local generation and editing on the shared RTX 3080",
    ),
}
MAX_INPUT_BYTES = 25 << 20

IMAGE_GENERATE_PARAMETERS = {
    "type": "object",
    "properties": {
        "prompt": {
            "type": "string",
            "description": "Detailed description of the image or requested edit.",
        },
        "aspect_ratio": {
            "type": "string",
            "enum": ["landscape", "square", "portrait"],
            "default": DEFAULT_ASPECT_RATIO,
            "description": "Requested output aspect ratio.",
        },
        "image_url": {
            "type": "string",
            "description": "Optional uploaded image URL or private local path to edit.",
        },
        "reference_image_urls": {
            "type": "array",
            "items": {"type": "string"},
            "description": "Optional reference images for the generation or edit.",
        },
    },
    "required": ["prompt"],
}

LOCAL_IMAGE_SCHEMA = {
    "name": "image_generate_local",
    "description": (
        "Generate or edit an image only with the private local FLUX model on "
        "Atlas hardware. Use this tool whenever the user says local, private, "
        "FLUX, on my hardware, or otherwise explicitly rejects a hosted image "
        "provider. Never substitute the generic image_generate tool for an "
        "explicit local request. The backend is already provisioned; never ask "
        "for an endpoint or tell the user to install Diffusers or ComfyUI."
    ),
    "parameters": IMAGE_GENERATE_PARAMETERS,
}

HOSTED_IMAGE_SCHEMA = {
    "name": "image_generate_hosted",
    "description": (
        "Generate or edit an image only with hosted OpenAI GPT Image at the "
        "highest configured quality. Use this tool when the user explicitly "
        "asks for OpenAI, GPT Image, or hosted image generation. Do not use it "
        "when the user explicitly requests local or private generation."
    ),
    "parameters": IMAGE_GENERATE_PARAMETERS,
}


def _broker_key() -> str:
    """Load the internal relay key from process env or the private .env."""
    value = os.environ.get("HERMES_IMAGE_BROKER_KEY", "").strip()
    if value:
        return value
    home = Path(os.environ.get("HERMES_HOME", "/opt/data"))
    try:
        for line in (home / ".env").read_text(encoding="utf-8").splitlines():
            if line.startswith("HERMES_IMAGE_BROKER_KEY="):
                return line.split("=", 1)[1].strip()
    except OSError:
        pass
    return ""


def _request(
    path: str, payload: dict[str, Any] | None, timeout: float
) -> dict[str, Any]:
    """Call the internal broker with a bounded authenticated JSON request."""
    key = _broker_key()
    if not key:
        raise RuntimeError("private image broker key is unavailable")
    data = None if payload is None else json.dumps(payload).encode("utf-8")
    request = Request(
        f"{BROKER_URL}{path}",
        data=data,
        method="GET" if data is None else "POST",
        headers={
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
    )
    try:
        with urlopen(request, timeout=timeout) as response:  # noqa: S310 - fixed internal URL
            return json.loads(response.read())
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")[:1000]
        raise RuntimeError(f"image broker returned HTTP {exc.code}: {body}") from exc
    except URLError as exc:
        raise RuntimeError(f"image broker is unavailable: {exc.reason}") from exc


def _local_image_data_url(value: str) -> str:
    """Convert a private uploaded image to a data URL for broker transport."""
    candidate = value.strip()
    lowered = candidate.lower()
    if lowered.startswith(("http://", "https://", "data:image/")):
        return candidate
    try:
        from agent.file_safety import get_read_block_error

        blocked = get_read_block_error(candidate)
        if blocked:
            raise ValueError(blocked)
    except ImportError:
        pass
    path = Path(os.path.expanduser(candidate)).resolve()
    raw = path.read_bytes()
    if not raw or len(raw) > MAX_INPUT_BYTES:
        raise ValueError("reference image must be between 1 byte and 25 MiB")
    if raw.startswith(b"\x89PNG\r\n\x1a\n"):
        mime = "image/png"
    elif raw.startswith(b"\xff\xd8\xff"):
        mime = "image/jpeg"
    elif raw.startswith((b"GIF87a", b"GIF89a")):
        mime = "image/gif"
    elif raw.startswith(b"RIFF") and raw[8:12] == b"WEBP":
        mime = "image/webp"
    else:
        raise ValueError("reference image must be PNG, JPEG, GIF, or WebP")
    return f"data:{mime};base64,{base64.b64encode(raw).decode('ascii')}"


class AtlasBrokerImageProvider(ImageGenProvider):
    """High-quality GPT Image generation through the owner-only broker."""

    @property
    def name(self) -> str:
        return "atlas-broker"

    @property
    def display_name(self) -> str:
        return "Atlas private image studio"

    def is_available(self) -> bool:
        try:
            result = _request("/health", None, 5.0)
            return bool(result.get("success") and result.get("available"))
        except Exception:
            return False

    def list_models(self) -> list[dict[str, Any]]:
        models = []
        for model, values in MODELS.items():
            display, _quality, *strengths = values
            local = model == "flux-2-klein-4b-local"
            models.append(
                {
                    "id": model,
                    "display": display,
                    "speed": "~2–15min" if local else "~30s–3min",
                    "strengths": strengths[0]
                    if strengths
                    else "OpenAI GPT Image 2 quality",
                    "price": "local compute" if local else "included account capacity",
                }
            )
        return models

    def default_model(self) -> str:
        return DEFAULT_MODEL

    def capabilities(self) -> dict[str, Any]:
        return {"modalities": ["text", "image"], "max_reference_images": 16}

    def generate(
        self,
        prompt: str,
        aspect_ratio: str = DEFAULT_ASPECT_RATIO,
        *,
        image_url: str | None = None,
        reference_image_urls: list[str] | None = None,
        **kwargs: Any,
    ) -> dict[str, Any]:
        prompt = (prompt or "").strip()
        aspect = resolve_aspect_ratio(aspect_ratio)
        if not prompt:
            return error_response(
                error="Prompt is required",
                error_type="invalid_argument",
                provider=self.name,
                aspect_ratio=aspect,
            )
        model = str(kwargs.get("model") or DEFAULT_MODEL)
        if model not in MODELS:
            model = DEFAULT_MODEL
        try:
            primary = _local_image_data_url(image_url) if image_url else None
            references = [
                _local_image_data_url(item)
                for item in (normalize_reference_images(reference_image_urls) or [])[
                    :16
                ]
            ]
            response = _request(
                "/v1/images/generations",
                {
                    "prompt": prompt,
                    "aspect_ratio": aspect,
                    "model": model,
                    "image_url": primary,
                    "reference_image_urls": references,
                },
                1200.0,
            )
            if not response.get("success"):
                return error_response(
                    error=str(response.get("error") or "image generation failed"),
                    error_type=str(response.get("error_type") or "provider_error"),
                    provider=self.name,
                    model=model,
                    prompt=prompt,
                    aspect_ratio=aspect,
                )
            image = save_b64_image(str(response["image_b64"]), prefix=f"atlas_{model}")
            return success_response(
                image=str(image),
                model=str(response.get("model") or model),
                prompt=prompt,
                aspect_ratio=aspect,
                provider=self.name,
                modality="image" if primary or references else "text",
                extra={
                    "quality": response.get("quality", MODELS[model][1]),
                    "size": response.get("size"),
                    "route": response.get("route"),
                    "requested_model": response.get("requested_model", model),
                    "hosted_fallback_reason": response.get("hosted_fallback_reason"),
                    "input_image_count": response.get(
                        "input_image_count", len(references) + bool(primary)
                    ),
                },
            )
        except Exception as exc:
            return error_response(
                error=f"Private image studio failed: {exc}",
                error_type="broker_error",
                provider=self.name,
                model=model,
                prompt=prompt,
                aspect_ratio=aspect,
            )


def _handle_image_generate(args: dict[str, Any], route: str) -> str:
    """Dispatch one image request through a fixed route."""
    prompt = str(args.get("prompt") or "").strip()
    if not prompt:
        return json.dumps(
            {
                "success": False,
                "image": None,
                "error": "prompt is required for image generation",
                "error_type": "invalid_argument",
            }
        )
    model = ROUTE_MODELS.get(route)
    if model is None:  # pragma: no cover - routes are fixed by the wrappers
        raise ValueError(f"unsupported image route: {route}")
    provider = AtlasBrokerImageProvider()
    result = provider.generate(
        prompt,
        str(args.get("aspect_ratio") or DEFAULT_ASPECT_RATIO),
        model=model,
        image_url=args.get("image_url"),
        reference_image_urls=args.get("reference_image_urls"),
    )
    if isinstance(result, dict):
        result.setdefault("requested_route", route)
    return json.dumps(result, ensure_ascii=False)


def _handle_local_image(args: dict[str, Any], **_kwargs: Any) -> str:
    """Generate an image without allowing hosted provider substitution."""
    return _handle_image_generate(args, "local")


def _handle_hosted_image(args: dict[str, Any], **_kwargs: Any) -> str:
    """Generate an image through the highest-quality hosted route."""
    return _handle_image_generate(args, "hosted")


def _image_tool_available() -> bool:
    """Expose the tool when at least one broker route is healthy."""
    return AtlasBrokerImageProvider().is_available()


def register(ctx: Any) -> None:
    """Register the broker and explicit local/hosted image tools."""
    ctx.register_image_gen_provider(AtlasBrokerImageProvider())
    ctx.register_tool(
        name="image_generate_local",
        toolset="image_gen",
        schema=LOCAL_IMAGE_SCHEMA,
        handler=_handle_local_image,
        check_fn=_image_tool_available,
        requires_env=[],
        is_async=False,
        description=LOCAL_IMAGE_SCHEMA["description"],
        emoji="🎨",
    )
    ctx.register_tool(
        name="image_generate_hosted",
        toolset="image_gen",
        schema=HOSTED_IMAGE_SCHEMA,
        handler=_handle_hosted_image,
        check_fn=_image_tool_available,
        requires_env=[],
        is_async=False,
        description=HOSTED_IMAGE_SCHEMA["description"],
        emoji="🎨",
    )