import json
import logging
import os
from collections.abc import AsyncIterator
from typing import Any

import httpx
from dotenv import load_dotenv
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
logger = logging.getLogger("mimo-gateway")

app = FastAPI(title="MiMo Local Gateway")

MODEL_MAP = {
    "claude-sonnet-4-5": "mimo-v2.5-pro",
    "anthropic/claude-sonnet-4-5": "mimo-v2.5-pro",
    "claude-sonnet-4-5-1m": "mimo-v2.5-pro[1m]",
    "anthropic/claude-sonnet-4-5-1m": "mimo-v2.5-pro[1m]",
}
DEFAULT_MODEL = "claude-sonnet-4-5"


def base_url() -> str:
    return os.getenv(
        "MIMO_BASE_URL", "https://api.xiaomimimo.com/anthropic"
    ).rstrip("/")


def api_key() -> str | None:
    return os.getenv("MIMO_API_KEY")


def remap_model(body: dict[str, Any]) -> str:
    incoming_model = body.get("model") or DEFAULT_MODEL
    mapped_model = MODEL_MAP.get(incoming_model, incoming_model)
    body["model"] = mapped_model
    return mapped_model


def upstream_headers(key: str) -> dict[str, str]:
    return {"api-key": key, "Content-Type": "application/json"}


async def read_json_body(request: Request) -> dict[str, Any] | JSONResponse:
    try:
        body = await request.json()
    except (json.JSONDecodeError, UnicodeDecodeError):
        return JSONResponse({"error": "Request body must be valid JSON"}, status_code=400)
    if not isinstance(body, dict):
        return JSONResponse({"error": "Request body must be a JSON object"}, status_code=400)
    return body


def response_from_upstream(response: httpx.Response) -> JSONResponse:
    try:
        content = response.json()
    except (json.JSONDecodeError, UnicodeDecodeError):
        content = {"raw": response.text}
    return JSONResponse(content=content, status_code=response.status_code)


@app.get("/v1/models")
async def list_models() -> dict[str, Any]:
    logger.info("path=/v1/models")
    return {
        "object": "list",
        "data": [
            {
                "id": "claude-sonnet-4-5",
                "object": "model",
                "display_name": "MiMo 2.5 Pro",
            },
            {
                "id": "claude-sonnet-4-5-1m",
                "object": "model",
                "display_name": "MiMo 2.5 Pro 1M",
            },
        ],
    }


@app.get("/health")
async def health() -> dict[str, Any]:
    logger.info("path=/health")
    return {"ok": True, "gateway": "mimo-local-gateway", "upstream": base_url()}


@app.post("/v1/messages")
async def messages(request: Request):
    key = api_key()
    if not key:
        return JSONResponse(
            {"error": "Missing MIMO_API_KEY environment variable"}, status_code=500
        )

    body = await read_json_body(request)
    if isinstance(body, JSONResponse):
        return body
    mapped_model = remap_model(body)
    logger.info("path=/v1/messages mapped_model=%s", mapped_model)
    url = f"{base_url()}/v1/messages"

    if body.get("stream") is True:
        client = httpx.AsyncClient(timeout=None)
        try:
            upstream = await client.send(
                client.build_request("POST", url, headers=upstream_headers(key), json=body),
                stream=True,
            )
        except httpx.RequestError as exc:
            await client.aclose()
            logger.warning("path=/v1/messages upstream_error=%s", type(exc).__name__)
            return JSONResponse({"error": "Unable to reach MiMo upstream"}, status_code=502)

        logger.info(
            "path=/v1/messages mapped_model=%s upstream_status=%s",
            mapped_model,
            upstream.status_code,
        )

        async def chunks() -> AsyncIterator[bytes]:
            try:
                async for chunk in upstream.aiter_bytes():
                    yield chunk
            finally:
                await upstream.aclose()
                await client.aclose()

        return StreamingResponse(
            chunks(), status_code=upstream.status_code, media_type="text/event-stream"
        )

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            upstream = await client.post(url, headers=upstream_headers(key), json=body)
    except httpx.RequestError as exc:
        logger.warning("path=/v1/messages upstream_error=%s", type(exc).__name__)
        return JSONResponse({"error": "Unable to reach MiMo upstream"}, status_code=502)

    logger.info(
        "path=/v1/messages mapped_model=%s upstream_status=%s",
        mapped_model,
        upstream.status_code,
    )
    return response_from_upstream(upstream)


@app.post("/v1/messages/count_tokens")
async def count_tokens(request: Request):
    body = await read_json_body(request)
    if isinstance(body, JSONResponse):
        return body
    mapped_model = remap_model(body)
    logger.info("path=/v1/messages/count_tokens mapped_model=%s", mapped_model)

    def fallback() -> JSONResponse:
        text = json.dumps(body, ensure_ascii=False, separators=(",", ":"))
        return JSONResponse({"input_tokens": max(1, len(text) // 4)})

    key = api_key()
    if not key:
        logger.warning("path=/v1/messages/count_tokens missing_api_key using_fallback=true")
        return fallback()

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            upstream = await client.post(
                f"{base_url()}/v1/messages/count_tokens",
                headers=upstream_headers(key),
                json=body,
            )
    except httpx.RequestError as exc:
        logger.warning(
            "path=/v1/messages/count_tokens upstream_error=%s using_fallback=true",
            type(exc).__name__,
        )
        return fallback()

    logger.info(
        "path=/v1/messages/count_tokens mapped_model=%s upstream_status=%s",
        mapped_model,
        upstream.status_code,
    )
    if upstream.status_code == 404:
        return fallback()
    return response_from_upstream(upstream)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "mimo_gateway:app",
        host=os.getenv("HOST", "127.0.0.1"),
        port=int(os.getenv("PORT", "8787")),
    )