import base64
import mimetypes
import os
import urllib.request

from fastapi.responses import HTMLResponse
from gradio import Server
from huggingface_hub import InferenceClient

MODEL = "moonshotai/Kimi-K3"
PROVIDER = "together"

SYSTEM_PROMPT = (
    "You are Kimi K3, Moonshot AI's open-weight native multimodal agentic model. "
    "Be helpful, concise, and accurate."
)


def to_data_url(file_ref) -> str:
    """Convert a Gradio file reference (dict, FileData, or local path) to a base64 data URL."""
    # Unwrap FileData / dict
    path = None
    if isinstance(file_ref, dict):
        path = file_ref.get("path") or file_ref.get("orig_name")
    elif hasattr(file_ref, "path"):
        path = file_ref.path
    elif isinstance(file_ref, str):
        path = file_ref

    if not path:
        raise ValueError(f"Could not resolve file from: {file_ref!r}")

    # If path is a URL (Gradio may store remote files), fetch it
    if path.startswith("http://") or path.startswith("https://"):
        with urllib.request.urlopen(path) as resp:
            data = resp.read()
        mime = mimetypes.guess_type(path)[0] or "image/png"
    else:
        if not os.path.exists(path):
            raise FileNotFoundError(f"File not found: {path}")
        with open(path, "rb") as f:
            data = f.read()
        mime = mimetypes.guess_type(path)[0] or "image/png"

    b64 = base64.b64encode(data).decode()
    return f"data:{mime};base64,{b64}"


app = Server()


@app.api(stream_every=0.1)
def chat(message: dict, history: list) -> dict:
    """Chat with Kimi K3 (Moonshot AI) via HF Inference Providers. Supports text and images."""
    token = os.environ.get("HF_TOKEN")
    if not token:
        yield {"role": "assistant", "content": "HF_TOKEN is not configured on this Space. Set it in the Space settings → Secrets."}
        return

    client = InferenceClient(api_key=token, provider=PROVIDER, timeout=600)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in history:
        if msg["role"] == "assistant" and isinstance(msg.get("content"), str):
            content = msg["content"]
            if "</think>" in content:
                content = content.split("</think>", 1)[1]
            messages.append({"role": "assistant", "content": content.strip()})
        elif msg["role"] == "user":
            messages.append({"role": "user", "content": msg["content"]})

    content = []
    for path in message.get("files", []):
        content.append({"type": "image_url", "image_url": {"url": to_data_url(path)}})
    if message.get("text"):
        content.append({"type": "text", "text": message["text"]})
    messages.append({"role": "user", "content": content if len(content) > 1 else (message.get("text") or "")})

    stream = client.chat.completions.create(
        model=MODEL,
        messages=messages,
        max_tokens=8192,
        stream=True,
    )

    reasoning, answer = "", ""
    for chunk in stream:
        if not chunk.choices:
            continue
        delta = chunk.choices[0].delta
        r = getattr(delta, "reasoning_content", None)
        if r:
            reasoning += r
        if delta.content:
            answer += delta.content
        out = ""
        if reasoning:
            out += f"<think>{reasoning}</think>"
        out += answer
        yield {"role": "assistant", "content": out}


@app.get("/", response_class=HTMLResponse)
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    with open(html_path, "r", encoding="utf-8") as f:
        return f.read()


app.launch(show_error=True)