import json
import os
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import httpx
import uvicorn
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, Response, StreamingResponse
from starlette.background import BackgroundTask


UPSTREAM_URL = os.environ.get("LITELLM_UPSTREAM_URL", "http://127.0.0.1:7861").rstrip("/")
MODEL_CATALOG_PATH = Path(os.environ.get("MODEL_CATALOG_PATH", "/app/config/model-catalog.json"))
USABLE_MODELS_PATH = Path(os.environ.get("USABLE_MODELS_PATH", "/app/config/usable-models.json"))
GENLABS_BASE_URL = os.environ.get("GENLABS_API_BASE", "https://api.genlabs.dev/deca/v1").rstrip("/")
GENLABS_MODELS = (
    "genlabs/deca-2.5-mini",
    "genlabs/deca-2.5-pro",
    "genlabs/deca-2.5-ultra",
)
NONSTREAM_UPSTREAM_MODELS = {
    model.strip()
    for model in os.environ.get("NONSTREAM_UPSTREAM_MODELS", "voidai/gpt-oss-120b").split(",")
    if model.strip()
}
CLOUDFLARE_IMAGE_MODELS = {
    model.strip()
    for model in os.environ.get(
        "CLOUDFLARE_IMAGE_MODELS",
        ",".join(
            [
                "cloudflare/@cf/black-forest-labs/flux-1-schnell",
                "cloudflare/@cf/black-forest-labs/flux-2-dev",
                "cloudflare/@cf/black-forest-labs/flux-2-klein-4b",
                "cloudflare/@cf/black-forest-labs/flux-2-klein-9b",
                "cloudflare/@cf/bytedance/stable-diffusion-xl-lightning",
                "cloudflare/@cf/leonardo/lucid-origin",
                "cloudflare/@cf/leonardo/phoenix-1.0",
                "cloudflare/@cf/lykon/dreamshaper-8-lcm",
                "cloudflare/@cf/runwayml/stable-diffusion-v1-5-img2img",
                "cloudflare/@cf/runwayml/stable-diffusion-v1-5-inpainting",
                "cloudflare/@cf/stabilityai/stable-diffusion-xl-base-1.0",
            ]
        ),
    ).split(",")
    if model.strip()
}
CLOUDFLARE_IMAGE_MAX_N = max(1, int(os.environ.get("CLOUDFLARE_IMAGE_MAX_N", "1")))
POLLINATIONS_IMAGE_MODELS = {
    model.strip()
    for model in os.environ.get(
        "POLLINATIONS_IMAGE_MODELS",
        "pollinations/flux,pollinations/zimage,pollinations/gptimage",
    ).split(",")
    if model.strip()
}
MODELSLAB_IMAGE_MODELS = {
    "modelslab/midjourney",
    "modelslab/anything-v3",
    "modelslab/wifu-diffusion",
    "modelslab/arcane-diffusion",
}
HOP_BY_HOP_HEADERS = {
    "connection",
    "content-length",
    "host",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailer",
    "transfer-encoding",
    "upgrade",
}

app = FastAPI()
client = httpx.AsyncClient(timeout=httpx.Timeout(300.0, connect=30.0))
catalog_metadata_cache: dict[str, dict[str, Any]] | None = None
usable_metadata_cache: dict[str, dict[str, Any]] | None = None


@app.on_event("startup")
async def startup_event():
    """Skip re-rendering; start-litellm.sh already wrote a static config."""
    pass


@app.on_event("shutdown")
async def shutdown_event() -> None:
    await client.aclose()


_render_config_done = False


def _try_render_config():
    """Try to render litellm config. Called on startup and first request."""
    global _render_config_done
    if _render_config_done:
        return
    import subprocess
    config_template = os.environ.get("LITELLM_CONFIG_TEMPLATE", "/app/config/config.yaml")
    config_rendered = os.environ.get("LITELLM_RENDERED_CONFIG", "/tmp/litellm-config.yaml")
    try:
        result = subprocess.run(
            ["python", "/app/scripts/render-config.py", config_template, config_rendered, "--include-legacy-aliases"],
            timeout=30,
            capture_output=True,
            text=True,
        )
        if result.returncode == 0:
            _render_config_done = True
        else:
            pass  # Will retry on next request
    except Exception:
        pass  # Will retry on next request


def clean_headers(headers: Any) -> dict[str, str]:
    return {
        key: value
        for key, value in dict(headers).items()
        if key.lower() not in HOP_BY_HOP_HEADERS
    }


def model_looks_free(model_id: str, name: str = "") -> bool:
    text = f"{model_id} {name}".lower()
    return (
        model_id.lower().endswith(":free")
        or ":free" in model_id.lower()
        or " free" in text
        or text.startswith("free ")
        or text.endswith("(free)")
    )


def pricing_is_free(pricing: Any) -> bool:
    if not isinstance(pricing, dict) or not pricing:
        return False
    numeric_prices = []
    for value in pricing.values():
        try:
            numeric_prices.append(float(value))
        except (TypeError, ValueError):
            pass
    return bool(numeric_prices) and all(price == 0 for price in numeric_prices)


def capability_from_mode(mode: str | None) -> list[str]:
    if not mode:
        return []
    normalized = mode.strip().lower().replace("-", "_")
    if normalized in {"image", "image_generation", "text_to_image"}:
        return ["image"]
    if normalized in {"audio_transcription", "transcription"}:
        return ["audio", "transcription"]
    if normalized in {"audio_speech", "speech", "text_to_speech"}:
        return ["audio", "speech"]
    if normalized in {"embedding", "embeddings"}:
        return ["embedding"]
    if normalized in {"rerank", "reranking"}:
        return ["rerank"]
    if normalized in {"chat", "completion", "responses"}:
        return ["text"]
    return []


def unique_strings(values: list[Any]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for value in values:
        if not isinstance(value, str):
            continue
        label = value.strip().lower().replace("-", "_")
        if not label or label in seen:
            continue
        seen.add(label)
        out.append(label)
    return out


def suffix_parts(suffix: Any) -> tuple[str, str, dict[str, Any]]:
    if isinstance(suffix, dict):
        alias = str(suffix.get("alias") or suffix.get("id") or suffix.get("model") or "")
        model = str(suffix.get("model") or alias)
        return alias, model, suffix
    value = str(suffix)
    return value, value, {}


def explicit_bool(value: Any) -> bool | None:
    if isinstance(value, bool):
        return value
    if isinstance(value, str):
        lowered = value.strip().lower()
        if lowered in {"true", "yes", "y", "1", "free"}:
            return True
        if lowered in {"false", "no", "n", "0", "paid"}:
            return False
    return None


def catalog_entry_metadata(group: dict[str, Any], suffix: Any) -> tuple[str, dict[str, Any]] | None:
    alias_prefix = str(group.get("alias_prefix") or "").strip()
    model_prefix = str(group.get("model_prefix") or "").strip()
    alias_suffix, model_suffix, suffix_meta = suffix_parts(suffix)
    if not alias_prefix or not alias_suffix:
        return None

    model_id = f"{alias_prefix}/{alias_suffix}"
    upstream_model = f"{model_prefix}/{model_suffix}" if model_prefix else model_suffix
    group_info = group.get("model_info") if isinstance(group.get("model_info"), dict) else {}
    suffix_info = suffix_meta.get("model_info") if isinstance(suffix_meta.get("model_info"), dict) else {}
    model_info = {**group_info, **suffix_info}
    mode = (
        suffix_meta.get("mode")
        or group.get("mode")
        or model_info.get("mode")
        or suffix_meta.get("task")
        or group.get("task")
    )
    mode = str(mode).strip() if mode else None
    raw_capabilities = []
    for source in (group, suffix_meta, model_info):
        capabilities = source.get("capabilities") if isinstance(source, dict) else None
        if isinstance(capabilities, list):
            raw_capabilities.extend(capabilities)
        elif isinstance(capabilities, dict):
            raw_capabilities.extend(
                key for key, enabled in capabilities.items() if enabled
            )

    pricing = suffix_meta.get("pricing") or group.get("pricing") or model_info.get("pricing")
    free_value = (
        explicit_bool(suffix_meta.get("free"))
        if "free" in suffix_meta
        else explicit_bool(suffix_meta.get("is_free"))
    )
    if free_value is None:
        free_value = (
            explicit_bool(group.get("free"))
            if "free" in group
            else explicit_bool(group.get("is_free"))
        )
    if free_value is None:
        free_value = pricing_is_free(pricing) or model_looks_free(model_id, str(model_info.get("name") or ""))

    capabilities = unique_strings([*capability_from_mode(mode), *raw_capabilities])
    provider = str(group.get("provider") or alias_prefix).strip()
    metadata: dict[str, Any] = {
        "id": model_id,
        "provider": provider,
        "source_model": upstream_model,
        "free": bool(free_value),
        "is_free": bool(free_value),
        "capabilities": capabilities,
        "catalog_source": "config/model-catalog.json",
    }
    if mode:
        metadata["task"] = mode
        metadata["mode"] = mode
    if isinstance(pricing, dict):
        metadata["pricing"] = pricing
    if model_info:
        metadata["model_info"] = {
            **model_info,
            "mode": mode or model_info.get("mode"),
            "capabilities": capabilities or model_info.get("capabilities", []),
            "free": bool(free_value),
            "is_free": bool(free_value),
        }
    else:
        metadata["model_info"] = {
            "mode": mode,
            "capabilities": capabilities,
            "free": bool(free_value),
            "is_free": bool(free_value),
        }
    return model_id, metadata


def merge_catalog_metadata(existing: dict[str, Any], incoming: dict[str, Any]) -> dict[str, Any]:
    merged = dict(existing)
    for key, value in incoming.items():
        if key == "model_info":
            current_info = merged.get("model_info")
            merged_info = dict(current_info) if isinstance(current_info, dict) else {}
            if isinstance(value, dict):
                for info_key, info_value in value.items():
                    current_value = merged_info.get(info_key)
                    if current_value in (None, [], {}) and info_value not in (None, [], {}):
                        merged_info[info_key] = info_value
                    elif (
                        info_key == "capabilities"
                        and isinstance(info_value, list)
                        and isinstance(current_value, list)
                        and len(info_value) > len(current_value)
                    ):
                        merged_info[info_key] = info_value
            if merged_info:
                merged["model_info"] = merged_info
            continue

        current = merged.get(key)
        if current in (None, [], {}) and value not in (None, [], {}):
            merged[key] = value
        elif (
            key == "capabilities"
            and isinstance(value, list)
            and isinstance(current, list)
            and len(value) > len(current)
        ):
            merged[key] = value

    return merged


def load_catalog_metadata() -> dict[str, dict[str, Any]]:
    """Return empty metadata — model-catalog.json is huge (3579 models) and
    parsing it on cpu-basic hangs. render-config already filtered the config."""
    global catalog_metadata_cache
    if catalog_metadata_cache is not None:
        return catalog_metadata_cache
    catalog_metadata_cache = {}
    return catalog_metadata_cache


def _load_probe_file(path_or_url: str, local_path: Path) -> dict | None:
    import urllib.request
    payload = None
    try:
        resp = urllib.request.urlopen(path_or_url, timeout=10)
        payload = json.loads(resp.read())
    except Exception:
        pass
    if payload is None:
        try:
            payload = json.loads(local_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            pass
    return payload if isinstance(payload, dict) else None


def _merge_probe(metadata: dict, payload: dict, flag: str, source: str) -> None:
    checked_at = payload.get("checked_at") or payload.get("image_checked_at") or payload.get("vision_checked_at")
    for model_id in payload.get("usable_model_ids", []) or payload.get("image_usable_model_ids", []) or payload.get("vision_usable_model_ids", []):
        if not isinstance(model_id, str) or not model_id.strip():
            continue
        item = metadata.setdefault(model_id.strip(), {})
        item[flag] = True
        item[f"verified_{flag}"] = True
        item[f"{flag}_checked_at"] = checked_at
        item[f"{flag}_source"] = source
        if flag == "chat_usable":
            item["usable"] = True
    for model in payload.get("models", []) or payload.get("image_models", []) or payload.get("vision_models", []):
        if not isinstance(model, dict):
            continue
        model_id = model.get("id")
        if not isinstance(model_id, str) or not model_id.strip():
            continue
        key = model_id.strip()
        item = metadata.setdefault(key, {flag: True, f"verified_{flag}": True, f"{flag}_checked_at": checked_at, f"{flag}_source": source})
        item[flag] = True
        item[f"{flag}_latency_ms"] = model.get("latency_ms")
        item[f"{flag}_http_status"] = model.get("http_status")
        if flag == "chat_usable":
            item["usable"] = True


def load_usable_metadata() -> dict[str, dict[str, Any]]:
    global usable_metadata_cache

    metadata: dict[str, dict[str, Any]] = {}

    chat_dataset_url = os.environ.get(
        "USABLE_MODELS_DATASET_URL",
        "https://huggingface.co/datasets/alchoholpad/litellm-usable-models/resolve/main/usable-models.json",
    )
    chat_payload = _load_probe_file(chat_dataset_url, USABLE_MODELS_PATH)
    if chat_payload:
        _merge_probe(metadata, chat_payload, "chat_usable", "config/usable-models.json")

    image_payload = _load_probe_file(
        "https://huggingface.co/datasets/alchoholpad/litellm-usable-models/resolve/main/usable-image-models.json",
        Path("/app/config/usable-image-models.json"),
    )
    if image_payload:
        _merge_probe(metadata, image_payload, "image_usable", "config/usable-image-models.json")

    vision_payload = _load_probe_file(
        "https://huggingface.co/datasets/alchoholpad/litellm-usable-models/resolve/main/usable-vision-models.json",
        Path("/app/config/usable-vision-models.json"),
    )
    if vision_payload:
        _merge_probe(metadata, vision_payload, "vision_usable", "config/usable-vision-models.json")

    usable_metadata_cache = metadata
    return metadata


def model_metadata(model_id: str) -> dict[str, Any]:
    return {
        **load_catalog_metadata().get(model_id, {}),
        **load_usable_metadata().get(model_id, {}),
    }


def runtime_extra_model_ids() -> set[str]:
    model_ids: set[str] = set()
    if genlabs_api_key():
        model_ids.update(GENLABS_MODELS)
    if cloudflare_api_token() and cloudflare_account_id():
        model_ids.update(CLOUDFLARE_IMAGE_MODELS)
    if pollinations_api_key():
        model_ids.update(POLLINATIONS_IMAGE_MODELS)
    if stablehorde_api_key():
        model_ids.update(STABLEHORDE_MODELS)
    return model_ids


def merge_model_metadata(raw_model: dict[str, Any], metadata: dict[str, Any]) -> dict[str, Any]:
    enriched = dict(raw_model)
    for key in (
        "provider",
        "free",
        "is_free",
        "pricing",
        "task",
        "mode",
        "capabilities",
        "catalog_source",
        "usable",
        "verified_usable",
        "usable_checked_at",
        "usable_source",
        "usable_latency_ms",
        "usable_http_status",
        "image_usable",
        "verified_image_usable",
        "image_usable_checked_at",
        "image_usable_source",
        "image_usable_latency_ms",
        "image_usable_http_status",
        "image_usable_carried_from_previous",
        "image_usable_last_probe_error_class",
        "image_usable_last_probe_http_status",
        "image_usable_last_probe_checked_at",
    ):
        if key in metadata and key not in enriched:
            enriched[key] = metadata[key]

    raw_info = enriched.get("model_info")
    merged_info = dict(raw_info) if isinstance(raw_info, dict) else {}
    metadata_info = metadata.get("model_info")
    if isinstance(metadata_info, dict):
        for key, value in metadata_info.items():
            if key not in merged_info or merged_info.get(key) in (None, [], {}):
                merged_info[key] = value
    if merged_info:
        enriched["model_info"] = merged_info
    return enriched


def genlabs_api_key() -> str | None:
    value = os.environ.get("GENLABS_API_KEY", "").strip()
    return value or None


def cloudflare_api_token() -> str | None:
    value = os.environ.get("CLOUDFLARE_API_TOKEN", "").strip()
    return value or None


def cloudflare_account_id() -> str | None:
    value = os.environ.get("CLOUDFLARE_ACCOUNT_ID", "").strip()
    return value or None


def pollinations_api_key() -> str | None:
    for name in ("POLLINATIONS_API_KEY", "POLLINATIONS_API_KEY_1"):
        value = os.environ.get(name, "").strip()
        if value:
            return value
    return None


STABLEHORDE_MODELS = {"stablehorde/stablehorde"}


def stablehorde_api_key() -> str | None:
    for name in ("STABLEHORDE_API_KEY", "STABLEHORDE_API_KEY_1"):
        value = os.environ.get(name, "").strip()
        if value:
            return value
    return None


def is_genlabs_chat_payload(payload: Any) -> bool:
    return isinstance(payload, dict) and payload.get("model") in GENLABS_MODELS


def is_nonstream_upstream_payload(payload: Any) -> bool:
    return (
        isinstance(payload, dict)
        and bool(payload.get("stream"))
        and str(payload.get("model", "")) in NONSTREAM_UPSTREAM_MODELS
    )


def cloudflare_image_model_id(model: str) -> str | None:
    if model in CLOUDFLARE_IMAGE_MODELS:
        return model.split("/", 1)[1]
    if model.startswith("@cf/") and f"cloudflare/{model}" in CLOUDFLARE_IMAGE_MODELS:
        return model
    return None


def pollinations_image_model_name(model: str) -> str | None:
    if model in POLLINATIONS_IMAGE_MODELS:
        return model.split("/", 1)[1]
    return None


def modelslab_image_model_id(model: str) -> str | None:
    if model in MODELSLAB_IMAGE_MODELS:
        return model.split("/", 1)[1]
    return None


def is_image_only_model(model: str) -> bool:
    return bool(cloudflare_image_model_id(model) or pollinations_image_model_name(model) or modelslab_image_model_id(model))


def genlabs_payload(payload: dict[str, Any], *, stream: bool) -> dict[str, Any]:
    out = dict(payload)
    out["model"] = str(payload["model"]).split("/", 1)[1]
    out["stream"] = stream
    return out


def sse_payload(data: dict[str, Any] | str) -> bytes:
    if isinstance(data, str):
        return f"data: {data}\n\n".encode("utf-8")
    return f"data: {json.dumps(data, separators=(',', ':'))}\n\n".encode("utf-8")


def completion_to_sse(payload: dict[str, Any]) -> Any:
    choices = payload.get("choices")
    choice = choices[0] if isinstance(choices, list) and choices else {}
    message = choice.get("message") if isinstance(choice, dict) else {}
    content = message.get("content") if isinstance(message, dict) else ""
    if content is None:
        content = ""
    if not isinstance(content, str):
        content = json.dumps(content, ensure_ascii=False)

    response_id = payload.get("id") or f"chatcmpl-{uuid.uuid4().hex}"
    created = payload.get("created") or int(time.time())
    model = payload.get("model") or "unknown"
    finish_reason = choice.get("finish_reason") if isinstance(choice, dict) else None
    finish_reason = finish_reason or "stop"

    def event(
        delta: dict[str, Any],
        *,
        finish: str | None = None,
        usage: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        out: dict[str, Any] = {
            "id": response_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [
                {
                    "index": 0,
                    "delta": delta,
                    "finish_reason": finish,
                }
            ],
        }
        if usage is not None:
            out["usage"] = usage
        return out

    yield sse_payload(event({"role": "assistant"}))
    if content:
        yield sse_payload(event({"content": content}))
    yield sse_payload(event({}, finish=finish_reason))

    usage = payload.get("usage")
    if isinstance(usage, dict):
        yield sse_payload(
            {
                "id": response_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": model,
                "choices": [],
                "usage": usage,
            }
        )
    yield sse_payload("[DONE]")


async def nonstream_upstream_chat(request: Request, payload: dict[str, Any]) -> Response:
    upstream_payload = dict(payload)
    upstream_payload["stream"] = False
    upstream_response = await client.post(
        f"{UPSTREAM_URL}/v1/chat/completions",
        params=request.query_params,
        headers=clean_headers(request.headers),
        json=upstream_payload,
    )

    if upstream_response.status_code < 200 or upstream_response.status_code >= 300:
        return Response(
            upstream_response.content,
            status_code=upstream_response.status_code,
            headers=clean_headers(upstream_response.headers),
            media_type=upstream_response.headers.get("content-type"),
        )

    try:
        response_payload = upstream_response.json()
    except json.JSONDecodeError:
        return Response(
            upstream_response.content,
            status_code=upstream_response.status_code,
            headers=clean_headers(upstream_response.headers),
            media_type=upstream_response.headers.get("content-type"),
        )

    return StreamingResponse(
        completion_to_sse(response_payload),
        status_code=upstream_response.status_code,
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


async def genlabs_stream(payload: dict[str, Any]) -> Response:
    api_key = genlabs_api_key()
    if not api_key:
        return JSONResponse({"error": "GENLABS_API_KEY is not configured"}, status_code=503)

    request = client.build_request(
        "POST",
        f"{GENLABS_BASE_URL}/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Accept": "text/event-stream",
            "Content-Type": "application/json",
        },
        json=genlabs_payload(payload, stream=True),
    )
    response = await client.send(request, stream=True)
    return StreamingResponse(
        response.aiter_raw(),
        status_code=response.status_code,
        headers=clean_headers(response.headers),
        background=BackgroundTask(response.aclose),
    )


async def genlabs_completion(payload: dict[str, Any]) -> Response:
    api_key = genlabs_api_key()
    if not api_key:
        return JSONResponse({"error": "GENLABS_API_KEY is not configured"}, status_code=503)

    model_name = str(payload["model"])
    content_parts: list[str] = []
    finish_reason = "stop"
    usage: dict[str, Any] | None = None
    response_id = f"chatcmpl-{uuid.uuid4().hex}"

    async with client.stream(
        "POST",
        f"{GENLABS_BASE_URL}/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Accept": "text/event-stream",
            "Content-Type": "application/json",
        },
        json=genlabs_payload(payload, stream=True),
    ) as response:
        if response.status_code < 200 or response.status_code >= 300:
            body = await response.aread()
            return Response(
                body,
                status_code=response.status_code,
                headers=clean_headers(response.headers),
                media_type=response.headers.get("content-type"),
            )

        async for line in response.aiter_lines():
            if not line.startswith("data:"):
                continue
            data = line.removeprefix("data:").strip()
            if not data or data == "[DONE]":
                continue
            try:
                chunk = json.loads(data)
            except json.JSONDecodeError:
                continue
            response_id = chunk.get("id") or response_id
            if isinstance(chunk.get("usage"), dict):
                usage = chunk["usage"]
            choices = chunk.get("choices")
            if not isinstance(choices, list) or not choices:
                continue
            choice = choices[0]
            if not isinstance(choice, dict):
                continue
            if choice.get("finish_reason"):
                finish_reason = str(choice["finish_reason"])
            delta = choice.get("delta")
            if isinstance(delta, dict) and isinstance(delta.get("content"), str):
                content_parts.append(delta["content"])

    return JSONResponse(
        {
            "id": response_id,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model_name,
            "choices": [
                {
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": "".join(content_parts),
                    },
                    "finish_reason": finish_reason,
                }
            ],
            "usage": usage
            or {
                "prompt_tokens": 0,
                "completion_tokens": 0,
                "total_tokens": 0,
            },
        }
    )


def openai_image_data(image_base64: str, prompt: str, media_type: str = "image/jpeg") -> dict[str, str]:
    return {
        "b64_json": image_base64,
        "url": f"data:{media_type};base64,{image_base64}",
        "revised_prompt": prompt,
    }


def cloudflare_image_payload(payload: dict[str, Any]) -> dict[str, Any]:
    out: dict[str, Any] = {"prompt": payload.get("prompt")}
    for key in ("steps", "seed"):
        if payload.get(key) is not None:
            out[key] = payload[key]
    return out


async def cloudflare_images_generations(payload: dict[str, Any]) -> Response:
    model = str(payload.get("model") or "")
    model_id = cloudflare_image_model_id(model)
    if not model_id:
        return JSONResponse(
            {"error": f"Unsupported Cloudflare image model: {model or '(missing)'}"},
            status_code=400,
        )

    prompt = payload.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        return JSONResponse({"error": "Image generation prompt is required"}, status_code=400)

    api_token = cloudflare_api_token()
    account_id = cloudflare_account_id()
    if not api_token or not account_id:
        return JSONResponse(
            {"error": "CLOUDFLARE_API_TOKEN and CLOUDFLARE_ACCOUNT_ID must be configured"},
            status_code=503,
        )

    requested_n = payload.get("n") if isinstance(payload.get("n"), int) else 1
    image_count = min(max(1, requested_n), CLOUDFLARE_IMAGE_MAX_N)
    images: list[dict[str, str]] = []
    for _ in range(image_count):
        response = await client.post(
            f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run/{model_id}",
            headers={
                "Authorization": f"Bearer {api_token}",
                "Content-Type": "application/json",
            },
            json=cloudflare_image_payload(payload),
        )

        content_type = response.headers.get("content-type", "")
        if response.status_code < 200 or response.status_code >= 300:
            return Response(
                response.content,
                status_code=response.status_code,
                headers=clean_headers(response.headers),
                media_type=content_type,
            )

        if content_type.startswith("image/"):
            import base64

            media_type = content_type.split(";", 1)[0]
            images.append(openai_image_data(base64.b64encode(response.content).decode("ascii"), prompt, media_type))
            continue

        try:
            response_payload = response.json()
        except json.JSONDecodeError:
            return Response(
                response.content,
                status_code=response.status_code,
                headers=clean_headers(response.headers),
                media_type=content_type,
            )

        result = response_payload.get("result") if isinstance(response_payload, dict) else None
        image_base64 = result.get("image") if isinstance(result, dict) else None
        if not isinstance(image_base64, str) or not image_base64:
            return JSONResponse(
                {
                    "error": "Cloudflare image response did not include result.image",
                    "provider_response": response_payload,
                },
                status_code=502,
            )
        images.append(openai_image_data(image_base64, prompt))

    return JSONResponse(
        {
            "created": int(time.time()),
            "data": images,
        }
    )


async def modelslab_images_generations(payload: dict[str, Any]) -> Response:
    model = str(payload.get("model") or "")
    model_id = modelslab_image_model_id(model)
    if not model_id:
        return JSONResponse(
            {"error": f"Unsupported ModelSLAB image model: {model or '(missing)'}"},
            status_code=400,
        )

    prompt = payload.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        return JSONResponse({"error": "Image generation prompt is required"}, status_code=400)

    api_key = os.environ.get("MODELSLAB_API_KEY", "").strip()
    if not api_key:
        return JSONResponse(
            {"error": "MODELSLAB_API_KEY must be configured"},
            status_code=503,
        )

    requested_n = payload.get("n") if isinstance(payload.get("n"), int) else 1
    image_count = min(max(1, requested_n), 4)
    images: list[dict[str, str]] = []

    for _ in range(image_count):
        response = await client.post(
            "https://modelslab.com/api/v6/images/text2img",
            headers={"Content-Type": "application/json"},
            json={
                "key": api_key,
                "model_id": model_id,
                "prompt": prompt,
                "negative_prompt": payload.get("negative_prompt", ""),
                "width": payload.get("size", "1024x1024").split("x")[0] if isinstance(payload.get("size"), str) else 1024,
                "height": payload.get("size", "1024x1024").split("x")[1] if isinstance(payload.get("size"), str) else 1024,
                "samples": 1,
            },
        )

        if response.status_code < 200 or response.status_code >= 300:
            return Response(
                response.content,
                status_code=response.status_code,
                headers=clean_headers(response.headers),
                media_type=response.headers.get("content-type"),
            )

        try:
            response_payload = response.json()
        except json.JSONDecodeError:
            return JSONResponse({"error": "Invalid JSON response from ModelSLAB"}, status_code=502)

        if response_payload.get("status") == "error":
            return JSONResponse(
                {"error": response_payload.get("message", "ModelSLAB error")},
                status_code=400,
            )

        image_urls = response_payload.get("output", [])
        if not image_urls:
            return JSONResponse(
                {"error": "ModelSLAB did not return any images", "provider_response": response_payload},
                status_code=502,
            )

        import base64
        for url in image_urls[:1]:
            if isinstance(url, str) and url.startswith("http"):
                img_response = await client.get(url)
                if img_response.status_code == 200:
                    media_type = img_response.headers.get("content-type", "image/png").split(";")[0]
                    images.append(openai_image_data(base64.b64encode(img_response.content).decode("ascii"), prompt, media_type))

    return JSONResponse(
        {
            "created": int(time.time()),
            "data": images,
        }
    )


def pollinations_image_payload(payload: dict[str, Any], model_name: str) -> dict[str, Any]:
    out: dict[str, Any] = {
        "prompt": payload.get("prompt"),
        "model": model_name,
    }
    for key in ("n", "size", "quality", "response_format", "user", "image", "safe"):
        if payload.get(key) is not None:
            out[key] = payload[key]
    return out


async def pollinations_images_generations(payload: dict[str, Any]) -> Response:
    model = str(payload.get("model") or "")
    model_name = pollinations_image_model_name(model)
    if not model_name:
        return JSONResponse(
            {"error": f"Unsupported Pollinations image model: {model or '(missing)'}"},
            status_code=400,
        )

    prompt = payload.get("prompt")
    if not isinstance(prompt, str) or not prompt.strip():
        return JSONResponse({"error": "Image generation prompt is required"}, status_code=400)

    api_key = pollinations_api_key()
    if not api_key:
        return JSONResponse({"error": "POLLINATIONS_API_KEY must be configured"}, status_code=503)

    response = await client.post(
        "https://gen.pollinations.ai/v1/images/generations",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json=pollinations_image_payload(payload, model_name),
    )
    return Response(
        response.content,
        status_code=response.status_code,
        headers=clean_headers(response.headers),
        media_type=response.headers.get("content-type"),
    )


async def genlabs_chat_completions(request: Request) -> Response:
    payload = await request.json()
    if bool(payload.get("stream")):
        return await genlabs_stream(payload)
    return await genlabs_completion(payload)


async def models_response(request: Request) -> Response:
    try:
        upstream = await client.get(
            f"{UPSTREAM_URL}/v1/models",
            params=request.query_params,
            headers=clean_headers(request.headers),
            timeout=5,
        )
    except Exception:
        # upstream litellm not ready yet; return static model list
        static_ids = ["groq/llama-3.3-70b-versatile", "sambanova/Meta-Llama-3.3-70B-Instruct", "gemini-flash-3"]
        return JSONResponse({
            "object": "list",
            "data": [{"id": mid, "object": "model", "created": 0, "owned_by": mid.split("/", 1)[0]} for mid in static_ids]
        })
    try:
        payload = upstream.json()
    except json.JSONDecodeError:
        return Response(
            upstream.content,
            status_code=upstream.status_code,
            headers=clean_headers(upstream.headers),
            media_type=upstream.headers.get("content-type"),
        )

    if isinstance(payload, dict) and isinstance(payload.get("data"), list):
        enriched_data: list[Any] = []
        existing = set()
        for item in payload["data"]:
            if isinstance(item, dict):
                model_id = item.get("id")
                if isinstance(model_id, str):
                    existing.add(model_id)
                    metadata = model_metadata(model_id)
                    enriched_data.append(merge_model_metadata(item, metadata) if metadata else item)
                    continue
            enriched_data.append(item)
        payload["data"] = enriched_data

        for model in sorted(runtime_extra_model_ids()):
            if model not in existing:
                metadata = model_metadata(model)
                payload["data"].append(
                    {
                        "id": model,
                        "object": "model",
                        "created": 0,
                        "owned_by": metadata.get("provider") or model.split("/", 1)[0],
                        **{key: value for key, value in metadata.items() if key != "id"},
                    }
                )
    return JSONResponse(payload, status_code=upstream.status_code)


def model_catalog_response() -> Response:
    metadata = {
        **load_catalog_metadata(),
    }
    for model_id, usable in load_usable_metadata().items():
        metadata[model_id] = {**metadata.get(model_id, {"id": model_id}), **usable}
    return JSONResponse(
        {
            "object": "model_catalog",
            "sources": [str(MODEL_CATALOG_PATH), str(USABLE_MODELS_PATH)],
            "count": len(metadata),
            "data": [metadata[key] for key in sorted(metadata)],
        }
    )


async def proxy_request(path: str, request: Request) -> Response:
    url = f"{UPSTREAM_URL}/{path}"
    body = await request.body()
    upstream_request = client.build_request(
        request.method,
        url,
        params=request.query_params,
        headers=clean_headers(request.headers),
        content=body,
    )
    upstream_response = await client.send(upstream_request, stream=True)
    return StreamingResponse(
        upstream_response.aiter_raw(),
        status_code=upstream_response.status_code,
        headers=clean_headers(upstream_response.headers),
        background=BackgroundTask(upstream_response.aclose),
    )


verify_lock = False
discover_lock = False
USABLE_MODELS_DATASET = os.environ.get("USABLE_MODELS_DATASET", "alchoholpad/litellm-usable-models")


@app.api_route("/admin/logs", methods=["GET"])
async def admin_logs() -> Response:
    try:
        log_path = Path("/tmp/verify.log")
        if log_path.exists():
            content = log_path.read_text(encoding="utf-8", errors="replace")
            return Response(content[-4096:], media_type="text/plain")
        return JSONResponse({"error": "no log file yet"}, status_code=404)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


@app.api_route("/admin/discover", methods=["POST"])
async def admin_discover() -> Response:
    global discover_lock
    if discover_lock:
        return JSONResponse({"status": "already running"}, status_code=409)
    discover_lock = True

    async def _run_discover():
        global discover_lock
        log_path = Path("/tmp/discover.log")
        log_path.write_text("starting catalog discovery...\n")
        try:
            catalog = json.loads(MODEL_CATALOG_PATH.read_text(encoding="utf-8"))
            groups = catalog.get("groups", [])
            discovered: dict[str, list[str]] = {}
            errors: dict[str, str] = {}

            for group in groups:
                alias = group.get("alias_prefix", "unknown")
                api_base = (group.get("params") or {}).get("api_base", "")
                key_env = group.get("api_key_env", "")
                if not api_base or not key_env:
                    continue
                api_key = os.environ.get(key_env, "")
                if not api_key:
                    errors[alias] = "no_api_key"
                    continue

                models_url = f"{api_base.rstrip('/')}/models"
                with open(log_path, "a") as log_f:
                    log_f.write(f"[{alias}] fetching {models_url}\n")
                try:
                    import httpx as _httpx
                    async with _httpx.AsyncClient(timeout=_httpx.Timeout(15.0)) as _client:
                        resp = await _client.get(models_url, headers={"Authorization": f"Bearer {api_key}"})
                        if resp.status_code >= 400:
                            errors[alias] = f"HTTP {resp.status_code}"
                            continue
                        data = resp.json()
                        model_list = data.get("data", data.get("models", data if isinstance(data, list) else []))
                        ids = [m.get("id", m) if isinstance(m, dict) else str(m) for m in model_list if m]
                        discovered[alias] = sorted(ids)
                        with open(log_path, "a") as log_f:
                            log_f.write(f"[{alias}] found {len(ids)} models\n")
                except Exception as exc:
                    errors[alias] = f"{type(exc).__name__}: {exc}"

            output = {
                "discovered_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
                "providers": {alias: {"count": len(models), "models": models} for alias, models in sorted(discovered.items())},
                "errors": errors,
                "total_models": sum(len(m) for m in discovered.values()),
                "total_providers": len(discovered),
            }
            out_path = Path("/app/config/discovered-models.json")
            tmp = out_path.with_suffix(".tmp")
            tmp.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n", encoding="utf-8")
            tmp.replace(out_path)
            with open(log_path, "a") as log_f:
                log_f.write(f"[done] {output['total_providers']} providers, {output['total_models']} models\n")
        except Exception as exc:
            with open(log_path, "a") as log_f:
                log_f.write(f"[fatal] {type(exc).__name__}: {exc}\n")
        finally:
            discover_lock = False

    import asyncio
    asyncio.get_event_loop().create_task(_run_discover())
    return JSONResponse({"status": "started"})


@app.api_route("/admin/discover/status", methods=["GET"])
async def admin_discover_status() -> Response:
    try:
        log_path = Path("/tmp/discover.log")
        if log_path.exists():
            return Response(log_path.read_text(encoding="utf-8", errors="replace")[-4096:], media_type="text/plain")
        return JSONResponse({"error": "no log file yet"}, status_code=404)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


@app.api_route("/admin/discover/results", methods=["GET"])
async def admin_discover_results() -> Response:
    try:
        results_path = Path("/app/config/discovered-models.json")
        if results_path.exists():
            return Response(results_path.read_text(encoding="utf-8", errors="replace"), media_type="application/json")
        return JSONResponse({"error": "no results yet, run /admin/discover first"}, status_code=404)
    except Exception as e:
        return JSONResponse({"error": str(e)}, status_code=500)


@app.api_route("/v1/usable-models", methods=["GET"])
async def usable_models_endpoint() -> Response:
    meta = load_catalog_metadata()
    usable = load_usable_metadata()
    chat, image, vision = [], [], []
    for model_id, info in usable.items():
        upstream = ""
        cat = meta.get(model_id)
        if cat and cat.get("source_model"):
            upstream = cat["source_model"]
        else:
            upstream = model_id
        if info.get("chat_usable"):
            chat.append(upstream)
        if info.get("image_usable"):
            image.append(upstream)
        if info.get("vision_usable"):
            vision.append(upstream)
    return JSONResponse({
        "chat": sorted(set(chat)),
        "image": sorted(set(image)),
        "vision": sorted(set(vision)),
    })


@app.api_route("/admin/verify", methods=["POST"])
async def admin_verify() -> Response:
    global verify_lock
    if verify_lock:
        return JSONResponse({"status": "already running"}, status_code=409)
    verify_lock = True

    async def _run_verify():
        global verify_lock, usable_metadata_cache
        log_path = Path("/tmp/verify.log")
        log_path.write_text("starting verify (chat + image + vision)...\n")
        try:
            import asyncio
            base_url_val = os.environ.get("LITELLM_BASE_URL") or f"https://{os.environ.get('SPACE_HOST', 'alchoholpad-litellm.hf.space')}"
            probes = [
                ("chat", ["--probe", "chat", "--output", "/app/config/usable-models.json"]),
                ("image", ["--probe", "image", "--image-scope", "all", "--no-filter", "--output", "/app/config/usable-image-models.json"]),
                ("vision", ["--probe", "vision", "--no-filter", "--retries", "3", "--retry-sleep", "5", "--concurrency", "4", "--output", "/app/config/usable-vision-models.json"]),
            ]
            for probe_idx, (probe_name, extra_args) in enumerate(probes):
                if probe_idx > 0:
                    cooldown = 300  # 5 min between probes to avoid rate limits
                    with open(log_path, "a") as log_f:
                        log_f.write(f"\n[cooldown] waiting {cooldown}s before {probe_name} probe...\n")
                    await asyncio.sleep(cooldown)
                with open(log_path, "a") as log_f:
                    log_f.write(f"\n=== {probe_name.upper()} PROBE ===\n")
                proc = await asyncio.create_subprocess_exec(
                    "python", "-u", "/app/scripts/verify-usable-models.py",
                    "--base-url", base_url_val,
                    "--concurrency", "8", "--timeout", "25", "--retries", "2", "--yes",
                    *extra_args,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.STDOUT,
                )
                with open(log_path, "a") as log_f:
                    async for line in proc.stdout:
                        log_f.write(line.decode(errors="replace"))
                        log_f.flush()
                await proc.wait()
                with open(log_path, "a") as log_f:
                    log_f.write(f"[{probe_name}] exit code: {proc.returncode}\n")
            # Upload all results to HF dataset
            try:
                from huggingface_hub import HfApi
                api = HfApi()
                hf_token = os.environ.get("HF_TOKEN")
                for fname in ["usable-models.json", "usable-image-models.json", "usable-vision-models.json"]:
                    fpath = Path(f"/app/config/{fname}")
                    if fpath.exists():
                        api.upload_file(
                            path_or_fileobj=str(fpath),
                            path_in_repo=fname,
                            repo_id=USABLE_MODELS_DATASET,
                            repo_type="dataset",
                            token=hf_token,
                        )
                usable_metadata_cache = None
                with open(log_path, "a") as log_f:
                    log_f.write("[upload] success\n")
            except Exception as e:
                with open(log_path, "a") as log_f:
                    log_f.write(f"[upload error] {e}\n")
        except Exception as e:
            with open(log_path, "a") as log_f:
                log_f.write(f"[fatal] {type(e).__name__}: {e}\n")
        finally:
            verify_lock = False

    import asyncio
    asyncio.get_event_loop().create_task(_run_verify())
    return JSONResponse({"status": "started"})


@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def route(path: str, request: Request) -> Response:
    normalized = f"/{path}"
    if request.method == "GET" and normalized in {"/v1/model-catalog", "/model-catalog"}:
        return model_catalog_response()
    if request.method == "GET" and normalized == "/v1/models":
        return await models_response(request)
    if request.method == "POST" and normalized == "/v1/chat/completions":
        _try_render_config()  # Ensure litellm config is rendered before first request
        try:
            payload = await request.json()
        except json.JSONDecodeError:
            return JSONResponse({"error": "Invalid JSON body"}, status_code=400)
        if is_image_only_model(str(payload.get("model") or "")):
            return JSONResponse(
                {
                    "error": {
                        "message": (
                            f"{payload.get('model')} is an image-only route in this wrapper. "
                            "Use /v1/images/generations for image generation and select a chat model "
                            "for normal Open WebUI conversations."
                        ),
                        "type": "invalid_request_error",
                        "code": "image_model_used_for_chat",
                    }
                },
                status_code=400,
            )
        if is_genlabs_chat_payload(payload):
            if bool(payload.get("stream")):
                return await genlabs_stream(payload)
            return await genlabs_completion(payload)
        if is_nonstream_upstream_payload(payload):
            return await nonstream_upstream_chat(request, payload)
    if request.method == "POST" and normalized in {"/v1/images/generations", "/images/generations"}:
        try:
            payload = await request.json()
        except json.JSONDecodeError:
            return JSONResponse({"error": "Invalid JSON body"}, status_code=400)
        if cloudflare_image_model_id(str(payload.get("model") or "")):
            return await cloudflare_images_generations(payload)
        if pollinations_image_model_name(str(payload.get("model") or "")):
            return await pollinations_images_generations(payload)
        if modelslab_image_model_id(str(payload.get("model") or "")):
            return await modelslab_images_generations(payload)

    return await proxy_request(path, request)


if __name__ == "__main__":
    uvicorn.run(
        app,
        host=os.environ.get("HOST", "0.0.0.0"),
        port=int(os.environ.get("PORT", "7860")),
    )

# build: 2026-06-21T10:49:48Z