#!/usr/bin/env python3
"""
HuggingFace Spaces → OpenAI-compatible API Proxy
=================================================
Exposes /v1/models and /v1/chat/completions (streaming + non-streaming).
Balances across multiple HF Spaces, queuing requests when all are busy.

Each space has a "type" that controls how the proxy talks to it:

  "openai"  — spaces that expose a real HTTP OpenAI-compatible API
               Health:  GET /health  →  {"ready": true/false, "status": "..."}
               Chat:    POST /v1/chat/completions  (streaming supported)
               Example: (none currently — all spaces use gradio type)

  "gradio"  — spaces built with Gradio, called via the gradio_client library
               so that requests are routed through the HF Pro GPU quota.
               Health:  GET /health  →  {"status": "ok", "model": "..."}
                        (no "ready" field — if it responds at all, it's ready)
               Chat:    gradio_client.Client(space_id, token=HF_TOKEN)
                          .predict(messages_json=..., api_name="/chat_completions")
               Token:   read from the HF_TOKEN environment variable / secret
               Example: qwen3-14b (fallback_module_trial spaces)
                        qwen3-30b-a3b (intelect_module spaces)
                        qwen3-coder-30b (coder_v2 spaces)
"""

import asyncio
import json
import logging
import os
import time
import uuid
import httpx
from gradio_client import Client as GradioClient

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional

# ─────────────────────────────────────────────────────────────────────────────
# CONFIGURE YOUR SPACES HERE
#
# Required fields for every space:
#   url       — base URL of the HF Space
#   model_id  — model name exposed to clients (e.g. Paperclip)
#   name      — human-readable label used in logs
#   type      — "openai" or "gradio" (controls how the proxy talks to it)
#
# Required for gradio spaces:
#   space_id  — HF repo id, e.g. "fomext/intelect_module_trial"
#               used by gradio_client so requests hit your Pro GPU quota
#
# Optional:
#   hf_token  — per-space HF token override (falls back to HF_TOKEN secret)
# ─────────────────────────────────────────────────────────────────────────────

SPACES = [

    # ── qwen3-14b  (gradio type — called via gradio_client) ─────────────────
    {
        "url":               "https://fomext-intelect-module-v3.hf.space",
        "space_id":          "fomext/intelect_module_v3",
        "model_id":          "qwen3-14b",
        "name":              "14b Reasoning (Space 5)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-intelect-module-v3-1.hf.space",
        "space_id":          "fomext/intelect_module_v3_1",
        "model_id":          "qwen3-14b",
        "name":              "14b Reasoning (Space 4)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-intelect-module-v3-2.hf.space",
        "space_id":          "fomext/intelect_module_v3_2",
        "model_id":          "qwen3-14b",
        "name":              "14b Reasoning (Space 3)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-intelect-module-v3-3.hf.space",
        "space_id":          "fomext/intelect_module_v3_3",
        "model_id":          "qwen3-14b",
        "name":              "14b Reasoning (Space 2)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-intelect-module-v3-4.hf.space",
        "space_id":          "fomext/intelect_module_v3_4",
        "model_id":          "qwen3-14b",
        "name":              "14b Reasoning (Space 1)",
        "type":              "gradio",
        "supports_thinking": False,
    },

    # ── qwen3-coder-30b  (gradio type — called via gradio_client) ───────────
    # NOTE: coder spaces do NOT accept the enable_thinking parameter
    {
        "url":               "https://fomext-coder-v2-trial.hf.space",
        "space_id":          "fomext/coder_v2_trial",
        "model_id":          "qwen3-coder-30b-a3b-instruct-fp8",
        "name":              "Coder 30b (Space 1)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-coder-v2-trial2.hf.space",
        "space_id":          "fomext/coder_v2_trial2",
        "model_id":          "qwen3-coder-30b-a3b-instruct-fp8",
        "name":              "Coder 30b (Space 2)",
        "type":              "gradio",
        "supports_thinking": False,
    },
    {
        "url":               "https://fomext-coder-v2-trial3.hf.space",
        "space_id":          "fomext/coder_v2_trial3",
        "model_id":          "qwen3-coder-30b-a3b-instruct-fp8",
        "name":              "Coder 30b (Space 3)",
        "type":              "gradio",
        "supports_thinking": False,
    },

    # ── qwen3-30b-a3b  (gradio type — called via gradio_client) ─────────────
    {
        "url":               "https://fomext-intelect_module_trial.hf.space",
        "space_id":          "fomext/intelect_module_trial",
        "model_id":          "qwen3-30b-a3b",
        "name":              "30b Reasoning (Space 1)",
        "type":              "gradio",
        "supports_thinking": True,
    },
    {
        "url":               "https://fomext-intelect_module_trial2.hf.space",
        "space_id":          "fomext/intelect_module_trial2",
        "model_id":          "qwen3-30b-a3b",
        "name":              "30b Reasoning (Space 2)",
        "type":              "gradio",
        "supports_thinking": True,
    },
    {
        "url":               "https://fomext-intelect_module_trial3.hf.space",
        "space_id":          "fomext/intelect_module_trial3",
        "model_id":          "qwen3-30b-a3b",
        "name":              "30b Reasoning (Space 3)",
        "type":              "gradio",
        "supports_thinking": True,
    },

]

# ── Model aliases ────────────────────────────────────────────────────────────
# Maps external model names (e.g. OpenAI names sent by Paperclip/OpenCode)
# to the actual model IDs configured in SPACES above.
# Add any new aliases here — no other code needs to change.
MODEL_ALIASES: dict[str, str] = {
    # OpenAI codex / GPT names → coder model
    "gpt-5.1-codex-mini":                  "qwen3-coder-30b-a3b-instruct-fp8",
    "gpt-5.1-codex":                       "qwen3-coder-30b-a3b-instruct-fp8",
    "code-davinci-002":                    "qwen3-coder-30b-a3b-instruct-fp8",
    # GPT-4-class names → 30b reasoning model
    "gpt-4o":                              "qwen3-30b-a3b",
    "gpt-4o-mini":                         "qwen3-14b",
    "gpt-4":                               "qwen3-30b-a3b",
    "gpt-4-turbo":                         "qwen3-14b",
    "gpt-4-turbo-preview":                 "qwen3-14b",
    # GPT-3.5 names → 14b model
    "gpt-3.5-turbo":                       "qwen3-14b",
    "gpt-3.5-turbo-16k":                   "qwen3-14b",
}

# Fallback model when the requested name isn't in SPACES or MODEL_ALIASES
DEFAULT_MODEL = "qwen3-14b"

# HF token for Gradio spaces — set this as a secret called HF_TOKEN
HF_TOKEN = os.environ.get("HF_TOKEN", "")


SPACE_READY_TIMEOUT = 600
# Seconds between health polls
POLL_INTERVAL = 10
# Upstream request timeout
REQUEST_TIMEOUT = 300

# ─────────────────────────────────────────────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)-8s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("hf-proxy")

app = FastAPI(title="HF Spaces OpenAI Proxy", version="2.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


# ── Space state ───────────────────────────────────────────────────────────────

class SpaceState:
    def __init__(self, cfg: dict):
        self.url:       str  = cfg["url"].rstrip("/")
        self.space_id:  str  = cfg.get("space_id", "")   # e.g. "fomext/intelect_module_trial"
        self.model_id:  str  = cfg["model_id"]
        self.name:      str  = cfg["name"]
        self.type:      str  = cfg["type"]          # "openai" | "gradio"
        self.hf_token:         str  = cfg.get("hf_token", "")
        self.supports_thinking: bool = cfg.get("supports_thinking", True)
        self.busy:             bool = False
        self.ready:     bool = False
        self.lock:           asyncio.Lock  = asyncio.Lock()
        self._ready_event:   asyncio.Event = asyncio.Event()

    def __repr__(self):
        s = "ready" if self.ready else "loading"
        b = "busy"  if self.busy  else "free"
        return f"<{self.name} [{self.type}] {s}/{b}>"


spaces: list[SpaceState] = [SpaceState(cfg) for cfg in SPACES]


# ── Health checks (type-aware) ────────────────────────────────────────────────

async def check_health_openai(space: SpaceState) -> bool:
    """openai spaces: GET /health must return {"ready": true}"""
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            r = await client.get(f"{space.url}/health")
            if r.status_code != 200:
                return False
            data = r.json()
            return bool(data.get("ready", False))
    except Exception:
        return False


async def check_health_gradio(space: SpaceState) -> bool:
    """
    Gradio spaces: GET /health returns {"status": "ok", "model": "..."}
    No "ready" field — if it responds with status=ok it IS ready.
    We also try the Gradio queue info endpoint as a fallback.
    """
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            r = await client.get(f"{space.url}/health")
            if r.status_code == 200:
                data = r.json()
                if data.get("status") == "ok":
                    return True
            # Fallback: Gradio exposes /info when the app is up
            r2 = await client.get(f"{space.url}/info")
            return r2.status_code == 200
    except Exception:
        return False


async def check_space_health(space: SpaceState) -> bool:
    if space.type == "openai":
        return await check_health_openai(space)
    else:
        return await check_health_gradio(space)


async def wait_until_ready(space: SpaceState):
    deadline = time.time() + SPACE_READY_TIMEOUT
    while time.time() < deadline:
        if await check_space_health(space):
            space.ready = True
            space._ready_event.set()
            log.info(f"Ready: {space}")
            return
        log.debug(f"Not ready yet: {space.name}")
        await asyncio.sleep(POLL_INTERVAL)
    log.warning(f"Timed out waiting for: {space.name}")


@app.on_event("startup")
async def startup():
    for space in spaces:
        asyncio.create_task(wait_until_ready(space))
    log.info(f"Proxy started — {len(spaces)} space(s) across "
             f"{len(set(s.model_id for s in spaces))} model(s)")


# ── Load balancer ─────────────────────────────────────────────────────────────

async def acquire_space(model_id: str) -> SpaceState:
    candidates = [s for s in spaces if s.model_id == model_id]
    if not candidates:
        raise HTTPException(404, detail=f"No space configured for model '{model_id}'")

    # Wait for at least one candidate to be ready
    ready_tasks = [asyncio.create_task(s._ready_event.wait()) for s in candidates]
    done, pending = await asyncio.wait(ready_tasks, return_when=asyncio.FIRST_COMPLETED)
    for t in pending:
        t.cancel()

    while True:
        for space in candidates:
            if space.ready and not space.busy:
                async with space.lock:
                    if not space.busy:
                        space.busy = True
                        log.info(f"Acquired {space.name}")
                        return space
        await asyncio.sleep(0.5)


def release_space(space: SpaceState):
    space.busy = False
    log.info(f"Released {space.name}")


# ── Chat adapters ─────────────────────────────────────────────────────────────
#
# openai spaces → forward body unchanged to /v1/chat/completions
# gradio spaces → call /run/chat_completions with messages serialised as JSON
#                 string; get back a plain text / JSON response and wrap it
#                 into an OpenAI-shaped reply for Paperclip.

async def call_openai_space(space: SpaceState, body: dict) -> dict:
    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        r = await client.post(
            f"{space.url}/v1/chat/completions",
            json=body,
            headers={"Content-Type": "application/json"},
        )
        r.raise_for_status()
        return r.json()


async def stream_openai_space(space: SpaceState, body: dict):
    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        async with client.stream(
            "POST",
            f"{space.url}/v1/chat/completions",
            json=body,
            headers={"Content-Type": "application/json"},
        ) as r:
            async for chunk in r.aiter_bytes():
                yield chunk


async def call_gradio_space(space: SpaceState, body: dict) -> dict:
    """
    Call a Gradio space via the gradio_client library so the request is
    routed through the caller's HF Pro GPU quota.

    gradio_client.Client.predict() is synchronous, so we run it in a
    thread-pool to avoid blocking the event loop.
    """
    messages        = body.get("messages", [])
    max_tokens      = body.get("max_tokens", 512)
    temperature     = body.get("temperature", 0.7)
    top_p           = body.get("top_p", 0.9)
    enable_thinking = body.get("enable_thinking", False)
    messages_json   = json.dumps(messages)

    # The upstream vLLM/transformers backend rejects temperature=0 with a
    # ValueError.  Clamp it to the smallest positive value that works.
    if temperature == 0:
        temperature = 0.01

    # Prefer per-space token, fall back to the global HF_TOKEN secret
    token = space.hf_token or HF_TOKEN or None

    # Use space_id (e.g. "fomext/intelect_module_trial") if set,
    # otherwise fall back to the bare URL.
    src = space.space_id if space.space_id else space.url

    def _call_sync() -> str:
        client = GradioClient(src, token=token)
        kwargs = dict(
            messages_json=messages_json,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            api_name="/chat_completions",
        )
        # Only pass enable_thinking to spaces that support it (e.g. reasoning
        # models). Coder spaces reject it with a keyword-argument error.
        if space.supports_thinking:
            kwargs["enable_thinking"] = enable_thinking
        return client.predict(**kwargs)

    loop = asyncio.get_event_loop()
    raw = await loop.run_in_executor(None, _call_sync)

    # raw is a JSON string returned by the Gradio endpoint
    if isinstance(raw, str):
        parsed = json.loads(raw)
    else:
        parsed = raw

    # If the space returned an error dict, surface it as a 502 rather than
    # silently wrapping the error string as model content.
    if "error" in parsed and "choices" not in parsed:
        raise HTTPException(502, detail=f"Upstream error: {parsed['error']}")

    if "choices" in parsed:
        return parsed

    content = parsed.get("content") or parsed.get("text") or str(parsed)
    return _wrap_as_openai(content, body.get("model", space.model_id))


def _wrap_as_openai(content: str, model_id: str) -> dict:
    """Wrap a plain text response into an OpenAI chat.completion shape."""
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:12]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": model_id,
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": content},
            "finish_reason": "stop",
        }],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


def _gradio_response_as_sse(openai_response: dict) -> bytes:
    """Convert a full OpenAI response dict into a single SSE event + DONE."""
    # Emit one delta chunk then [DONE]
    content = openai_response["choices"][0]["message"]["content"]
    chunk = {
        "id": openai_response["id"],
        "object": "chat.completion.chunk",
        "created": openai_response["created"],
        "model": openai_response["model"],
        "choices": [{
            "index": 0,
            "delta": {"role": "assistant", "content": content},
            "finish_reason": "stop",
        }],
    }
    data  = f"data: {json.dumps(chunk)}\n\n".encode()
    done  = b"data: [DONE]\n\n"
    return data + done


# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/")
async def root():
    return {"status": "ok", "spaces": len(spaces)}


@app.get("/health")
async def health():
    statuses = [
        {
            "name":  s.name,
            "model": s.model_id,
            "type":  s.type,
            "ready": s.ready,
            "busy":  s.busy,
        }
        for s in spaces
    ]
    return {
        "ready":  any(s.ready for s in spaces),
        "spaces": statuses,
    }


@app.get("/v1/models")
async def list_models():
    seen, models = set(), []
    for s in spaces:
        if s.model_id not in seen:
            seen.add(s.model_id)
            models.append({
                "id":       s.model_id,
                "object":   "model",
                "created":  0,
                "owned_by": "huggingface-spaces",
            })
    return {"object": "list", "data": models}


@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    body      = await request.json()
    model_id  = body.get("model", "")
    is_stream = body.get("stream", False)

    # Resolve any alias (e.g. "gpt-5.1-codex-mini" → "qwen3-coder-30b-a3b-instruct-fp8")
    # then fall back to DEFAULT_MODEL if the name is still unknown.
    resolved_id = MODEL_ALIASES.get(model_id, model_id) or DEFAULT_MODEL
    if resolved_id != model_id:
        log.info(f"Model alias: '{model_id}' → '{resolved_id}'")
        model_id = resolved_id
    if not any(s.model_id == model_id for s in spaces):
        log.warning(f"Unknown model '{model_id}', falling back to '{DEFAULT_MODEL}'")
        model_id = DEFAULT_MODEL
    body["model"] = model_id  # keep body in sync so upstream sees the real name

    space = await acquire_space(model_id)

    try:
        # ── openai-type space ─────────────────────────────────────────────
        if space.type == "openai":
            if is_stream:
                return StreamingResponse(
                    _stream_openai(space, body),
                    media_type="text/event-stream",
                )
            else:
                return await _non_stream_openai(space, body)

        # ── gradio-type space ─────────────────────────────────────────────
        else:
            # Gradio spaces don't support true streaming from this proxy.
            # We call the endpoint, get the full response, then either
            # return it directly or wrap it as a single SSE event.
            try:
                response = await call_gradio_space(space, body)
                release_space(space)
            except Exception as e:
                release_space(space)
                log.error(f"Gradio error ({space.name}): {e}")
                raise HTTPException(502, detail=f"Upstream error: {e}")

            if is_stream:
                # Paperclip asked for streaming — fake it with one big chunk
                sse_bytes = _gradio_response_as_sse(response)
                async def _single_chunk():
                    yield sse_bytes
                return StreamingResponse(_single_chunk(), media_type="text/event-stream")
            else:
                return JSONResponse(content=response)

    except HTTPException:
        raise
    except Exception:
        release_space(space)
        raise


async def _non_stream_openai(space: SpaceState, body: dict):
    try:
        result = await call_openai_space(space, body)
        release_space(space)
        return JSONResponse(content=result)
    except Exception as e:
        release_space(space)
        log.error(f"Upstream error ({space.name}): {e}")
        raise HTTPException(502, detail=f"Upstream error: {e}")


async def _stream_openai(space: SpaceState, body: dict):
    try:
        async for chunk in stream_openai_space(space, body):
            yield chunk
    except Exception as e:
        log.error(f"Stream error ({space.name}): {e}")
        yield b"data: [DONE]\n\n"
    finally:
        release_space(space)