import json
import os
import subprocess
import time
import uuid
from pathlib import Path
from threading import Lock, Thread
from typing import Any, Dict, Iterable, List

os.environ.setdefault("HF_HOME", "/tmp/hf_home")
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")
os.environ.setdefault("GRADIO_SSR_MODE", "false")
for _path in (os.environ["HF_HOME"], os.environ["HF_MODULES_CACHE"], os.environ["MPLCONFIGDIR"]):
    os.makedirs(_path, exist_ok=True)

import spaces  # noqa: E402

import httpx  # noqa: E402
import gradio as gr  # noqa: E402
from fastapi import Request  # noqa: E402
from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse  # noqa: E402
from starlette.background import BackgroundTask  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer  # noqa: E402
import torch  # noqa: E402


MODEL_ID = "WeiboAI/VibeThinker-3B"
MODEL_REVISION = "main"
DEFAULT_MAX_NEW_TOKENS = 4096
MAX_NEW_TOKENS = 4096
ROOT_DIR = Path(__file__).resolve().parent
CHAT_UI_DIR = ROOT_DIR / "chat-ui"
CHAT_UI_PORT = int(os.environ.get("CHAT_UI_PORT", "3000"))
CHAT_UI_URL = f"http://127.0.0.1:{CHAT_UI_PORT}"
CHAT_UI_BUILD = CHAT_UI_DIR / "build" / "index.js"
CHAT_UI_DB = Path(os.environ.get("CHAT_UI_DB", "/tmp/vibethinker-chat-ui-db"))

_model_lock = Lock()
_model = None
_tokenizer = None
_model_device = "cpu"
_chat_ui_lock = Lock()
_chat_ui_process: subprocess.Popen | None = None


def _download_model() -> None:
    print(f"Downloading {MODEL_ID} to the local Hub cache...", flush=True)
    snapshot_download(
        repo_id=MODEL_ID,
        revision=MODEL_REVISION,
        ignore_patterns=["*.msgpack", "*.h5", "*.ot", "*.onnx"],
    )
    print("Model files are present in the local Hub cache.", flush=True)


def _load_model_cpu() -> None:
    global _model, _tokenizer
    if _model is not None and _tokenizer is not None:
        return

    with _model_lock:
        if _model is not None and _tokenizer is not None:
            return

        _download_model()
        print("Loading tokenizer...", flush=True)
        _tokenizer = AutoTokenizer.from_pretrained(
            MODEL_ID,
            revision=MODEL_REVISION,
            trust_remote_code=True,
        )
        if _tokenizer.pad_token_id is None and _tokenizer.eos_token_id is not None:
            _tokenizer.pad_token = _tokenizer.eos_token

        print("Loading model on CPU...", flush=True)
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            revision=MODEL_REVISION,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
            attn_implementation="sdpa",
        ).eval()
        print("Model loaded on CPU.", flush=True)


def _ensure_model_on_cuda() -> None:
    global _model_device
    _load_model_cpu()
    if _model_device == "cuda":
        return
    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is not available. Confirm the Space is running on zero-a10g.")
    print("Moving model to CUDA...", flush=True)
    _model.to("cuda")
    _model_device = "cuda"
    print("Model is ready on CUDA.", flush=True)


def _clean_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, str]]:
    cleaned: List[Dict[str, str]] = []
    for item in messages:
        if not isinstance(item, dict):
            continue
        role = str(item.get("role", "")).strip()
        content = str(item.get("content", "")).strip()
        if role in {"system", "user", "assistant"} and content:
            cleaned.append({"role": role, "content": content})
    return cleaned


def _messages_from_json(history_json: str) -> List[Dict[str, str]]:
    try:
        raw = json.loads(history_json)
    except json.JSONDecodeError as exc:
        raise ValueError("history_json must be valid JSON") from exc
    if not isinstance(raw, list):
        raise ValueError("history_json must encode a list of chat messages")
    messages = _clean_messages(raw)
    if not messages or messages[-1]["role"] != "user":
        raise ValueError("The final message must be a user message")
    return messages


def _coerce_generation_args(
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
) -> Dict[str, Any]:
    max_new_tokens = max(1, min(int(max_new_tokens or DEFAULT_MAX_NEW_TOKENS), MAX_NEW_TOKENS))
    temperature = max(0.0, min(float(temperature), 2.0))
    top_p = max(0.05, min(float(top_p), 1.0))
    repetition_penalty = max(0.8, min(float(repetition_penalty), 1.5))
    return {
        "max_new_tokens": max_new_tokens,
        "do_sample": temperature > 0,
        "temperature": max(temperature, 1e-5),
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "pad_token_id": _tokenizer.pad_token_id,
        "eos_token_id": _tokenizer.eos_token_id,
    }


def _format_prompt(messages: List[Dict[str, str]]) -> str:
    return _tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )


def _estimate_duration(
    history_json: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
    *args: Any,
    **kwargs: Any,
) -> int:
    del history_json, temperature, top_p, repetition_penalty, args, kwargs
    return min(240, max(60, 40 + int(max_new_tokens or DEFAULT_MAX_NEW_TOKENS) // 12))


@spaces.GPU(duration=1)
def _zerogpu_probe() -> str:
    return "ready"


def _generate_stream(
    messages: List[Dict[str, str]],
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
) -> Iterable[str]:
    _ensure_model_on_cuda()
    prompt = _format_prompt(messages)
    inputs = _tokenizer([prompt], return_tensors="pt").to(_model.device)
    generation_args = _coerce_generation_args(
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
    )
    streamer = TextIteratorStreamer(
        _tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        timeout=180,
    )
    generation_kwargs = {
        **inputs,
        **generation_args,
        "streamer": streamer,
    }
    worker = Thread(target=_model.generate, kwargs=generation_kwargs, daemon=True)
    worker.start()

    partial = ""
    for token in streamer:
        partial += token
        yield partial
    worker.join(timeout=1)


@spaces.GPU(duration=_estimate_duration)
def _gpu_chat_stream(
    history_json: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    repetition_penalty: float,
) -> Iterable[str]:
    messages = _messages_from_json(history_json)
    yield from _generate_stream(
        messages=messages,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=repetition_penalty,
    )


app = gr.Server()


def _chat_ui_model_config() -> str:
    return json.dumps(
        [
            {
                "id": MODEL_ID,
                "name": MODEL_ID,
                "displayName": "VibeThinker-3B",
                "description": "Reasoning-focused 3B model hosted on ZeroGPU.",
                "modelUrl": f"https://huggingface.co/{MODEL_ID}",
                "parameters": {
                    "temperature": 1.0,
                    "top_p": 0.95,
                    "max_tokens": DEFAULT_MAX_NEW_TOKENS,
                },
                "supportsReasoning": True,
            }
        ],
        separators=(",", ":"),
    )


def _chat_ui_env() -> Dict[str, str]:
    env = os.environ.copy()
    env.update(
        {
            "OPENAI_BASE_URL": "http://127.0.0.1:7860/v1",
            "OPENAI_API_KEY": "sk-local",
            "USE_USER_TOKEN": "false",
            "AUTOMATIC_LOGIN": "false",
            "ALLOW_IFRAME": "true",
            "PUBLIC_APP_NAME": "VibeThinker",
            "PUBLIC_APP_ASSETS": "chatui",
            "PUBLIC_APP_DESCRIPTION": "VibeThinker-3B hosted on ZeroGPU.",
            "PUBLIC_ORIGIN": os.environ.get(
                "PUBLIC_ORIGIN",
                "https://mike0021-vibethinker-3b-zerogpu.hf.space",
            ),
            "MONGO_STORAGE_PATH": str(CHAT_UI_DB),
            "MONGODB_DB_NAME": "chat-ui",
            "MONGODB_DIRECT_CONNECTION": "false",
            "COOKIE_NAME": "vibethinker-chat-session",
            "HUSKY": "0",
            "ENABLE_CONFIG_MANAGER": "false",
            "LLM_SUMMARIZATION": "false",
            "TASK_MODEL": MODEL_ID,
            "MODELS": _chat_ui_model_config(),
            "PORT": str(CHAT_UI_PORT),
            "HOST": "127.0.0.1",
            "BODY_SIZE_LIMIT": os.environ.get("BODY_SIZE_LIMIT", "15728640"),
        }
    )
    return env


def _run_chat_ui_command(command: List[str], env: Dict[str, str]) -> None:
    print(f"[chat-ui] running: {' '.join(command)}", flush=True)
    subprocess.run(command, cwd=CHAT_UI_DIR, env=env, check=True)


def _ensure_chat_ui_build(env: Dict[str, str]) -> None:
    if not CHAT_UI_DIR.exists():
        raise RuntimeError("chat-ui source directory is missing from the Space.")

    node_modules = CHAT_UI_DIR / "node_modules"
    if not node_modules.exists():
        _run_chat_ui_command(["npm", "ci"], env)

    if not CHAT_UI_BUILD.exists():
        _run_chat_ui_command(["npm", "run", "build"], env)


def _chat_ui_is_ready() -> bool:
    try:
        with httpx.Client(timeout=2.0) as client:
            response = client.get(f"{CHAT_UI_URL}/healthcheck")
        return response.status_code < 500
    except Exception:
        return False


def _start_chat_ui() -> None:
    global _chat_ui_process
    with _chat_ui_lock:
        if _chat_ui_process is not None and _chat_ui_process.poll() is None:
            return

        env = _chat_ui_env()
        _ensure_chat_ui_build(env)
        print(f"[chat-ui] starting on {CHAT_UI_URL}", flush=True)
        _chat_ui_process = subprocess.Popen(
            [
                "node",
                "--dns-result-order=ipv4first",
                str(CHAT_UI_BUILD),
                "--",
                "--host",
                "127.0.0.1",
                "--port",
                str(CHAT_UI_PORT),
            ],
            cwd=CHAT_UI_DIR,
            env=env,
        )


def _start_chat_ui_background() -> None:
    try:
        time.sleep(3)
        _start_chat_ui()
    except Exception as exc:
        print(f"[chat-ui] failed to start: {exc}", flush=True)


async def _wait_for_chat_ui(timeout: float = 240.0) -> None:
    if _chat_ui_process is None or _chat_ui_process.poll() is not None:
        Thread(target=_start_chat_ui_background, daemon=True).start()

    deadline = time.time() + timeout
    while time.time() < deadline:
        if _chat_ui_is_ready():
            return
        await anyio_sleep(1.0)

    raise RuntimeError("chat-ui did not become ready before the timeout.")


async def anyio_sleep(seconds: float) -> None:
    import anyio

    await anyio.sleep(seconds)


@app.on_event("startup")
async def start_chat_ui_on_startup() -> None:
    Thread(target=_start_chat_ui_background, daemon=True).start()


@app.api(name="zerogpu_probe", concurrency_limit=1, time_limit=30)
def zerogpu_probe() -> str:
    return _zerogpu_probe()


@app.api(name="chat", concurrency_limit=1, time_limit=300, stream_every=0.25)
def chat(
    history_json: str,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 1.0,
    top_p: float = 0.95,
    repetition_penalty: float = 1.0,
) -> Iterable[str]:
    yield from _gpu_chat_stream(
        history_json,
        max_new_tokens,
        temperature,
        top_p,
        repetition_penalty,
    )


@app.get("/health")
async def health() -> Dict[str, str]:
    return {"status": "ok", "model": MODEL_ID}


@app.get("/v1/models")
async def openai_models() -> Dict[str, Any]:
    created = int(time.time())
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_ID,
                "object": "model",
                "created": created,
                "owned_by": "WeiboAI",
            }
        ],
    }


def _messages_for_openai(payload: Dict[str, Any]) -> List[Dict[str, str]]:
    raw_messages = payload.get("messages")
    if not isinstance(raw_messages, list):
        raise ValueError("messages must be a list")
    return _clean_messages(raw_messages)


@app.post("/v1/chat/completions", response_model=None)
async def openai_chat_completions(request: Request):
    try:
        payload = await request.json()
        messages = _messages_for_openai(payload)
        if not messages:
            raise ValueError("messages cannot be empty")
        max_new_tokens = int(payload.get("max_tokens") or payload.get("max_new_tokens") or DEFAULT_MAX_NEW_TOKENS)
        temperature = float(payload.get("temperature", 1.0))
        top_p = float(payload.get("top_p", 0.95))
        repetition_penalty = float(payload.get("repetition_penalty", 1.0))
        stream = bool(payload.get("stream", False))
    except Exception as exc:
        return JSONResponse({"error": {"message": str(exc), "type": "invalid_request_error"}}, status_code=400)

    completion_id = f"chatcmpl-{uuid.uuid4().hex}"
    created = int(time.time())

    if stream:
        def events() -> Iterable[bytes]:
            last_text = ""
            for text in _gpu_chat_stream(
                json.dumps(messages),
                max_new_tokens,
                temperature,
                top_p,
                repetition_penalty,
            ):
                delta = text[len(last_text) :]
                last_text = text
                chunk = {
                    "id": completion_id,
                    "object": "chat.completion.chunk",
                    "created": created,
                    "model": MODEL_ID,
                    "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}],
                }
                yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode("utf-8")
            final = {
                "id": completion_id,
                "object": "chat.completion.chunk",
                "created": created,
                "model": MODEL_ID,
                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
            }
            yield f"data: {json.dumps(final, ensure_ascii=False)}\n\n".encode("utf-8")
            yield b"data: [DONE]\n\n"

        return StreamingResponse(events(), media_type="text/event-stream")

    final_text = ""
    for final_text in _gpu_chat_stream(
        json.dumps(messages),
        max_new_tokens,
        temperature,
        top_p,
        repetition_penalty,
    ):
        pass
    return JSONResponse(
        {
            "id": completion_id,
            "object": "chat.completion",
            "created": created,
            "model": MODEL_ID,
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": final_text},
                    "finish_reason": "stop",
                }
            ],
        }
    )


@app.get("/gradio_api/startup-events")
async def gradio_startup_events() -> Dict[str, str]:
    return {"status": "ok"}


HOP_BY_HOP_HEADERS = {
    "connection",
    "keep-alive",
    "proxy-authenticate",
    "proxy-authorization",
    "te",
    "trailers",
    "transfer-encoding",
    "upgrade",
}


async def _close_upstream(response: httpx.Response, client: httpx.AsyncClient) -> None:
    await response.aclose()
    await client.aclose()


@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS", "HEAD"])
async def proxy_chat_ui(path: str, request: Request):
    host = request.headers.get("host", "")
    if path == "" and host.startswith(("localhost:", "127.0.0.1:", "0.0.0.0:")) and not _chat_ui_is_ready():
        return PlainTextResponse("chat-ui is starting", status_code=200)

    try:
        await _wait_for_chat_ui()
    except Exception as exc:
        return PlainTextResponse(f"chat-ui is starting or failed to start: {exc}", status_code=503)

    target_url = f"{CHAT_UI_URL}/{path}"
    if request.url.query:
        target_url = f"{target_url}?{request.url.query}"

    excluded_request_headers = HOP_BY_HOP_HEADERS | {"host", "content-length"}
    headers = {
        key: value
        for key, value in request.headers.items()
        if key.lower() not in excluded_request_headers
    }
    headers["x-forwarded-host"] = request.headers.get("host", "")
    headers["x-forwarded-proto"] = request.url.scheme

    client = httpx.AsyncClient(timeout=None, follow_redirects=False)
    try:
        upstream_request = client.build_request(
            request.method,
            target_url,
            headers=headers,
            content=await request.body(),
        )
        upstream_response = await client.send(upstream_request, stream=True)
    except Exception as exc:
        await client.aclose()
        return PlainTextResponse(f"chat-ui proxy error: {exc}", status_code=502)

    excluded_response_headers = HOP_BY_HOP_HEADERS | {"content-length"}
    response_headers = {
        key: value
        for key, value in upstream_response.headers.items()
        if key.lower() not in excluded_response_headers
    }
    return StreamingResponse(
        upstream_response.aiter_raw(),
        status_code=upstream_response.status_code,
        headers=response_headers,
        background=BackgroundTask(_close_upstream, upstream_response, client),
    )


demo = app


if __name__ == "__main__":
    app.launch()