"""
Nova-1 Chat — Powered by gradio.Server (Headless Gradio)
Custom ChatGPT-style frontend with ZeroGPU streaming support.

Requirements:
  pip install gradio>=6.0 transformers torch spaces

Important for HF Spaces:
  - Set environment variable GRADIO_SSR_MODE=false
  - The file index.html must be in the same directory as app.py
"""

import os
import traceback
from threading import Thread

import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

from gradio import Server
from fastapi.responses import HTMLResponse

# ─── Model Loading ────────────────────────────────────────────────────────────

model_name = "Smilyai-labs/Nova-1-Standard-Preview"

print("🔄 Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
print("✅ Tokenizer loaded!")

print("🔄 Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    dtype=torch.bfloat16,
    device_map="auto",
    low_cpu_mem_usage=True,
)
model.eval()
print("✅ Model loaded!")


# ─── Helpers ──────────────────────────────────────────────────────────────────

def format_chatml(messages, system_prompt: str) -> str:
    out = ""
    if system_prompt and system_prompt.strip():
        out += f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
    for m in messages:
        if m["role"] == "user":
            out += f"<|im_start|>user\n{m['content']}<|im_end|>\n"
        elif m["role"] == "assistant":
            out += f"<|im_start|>assistant\n{m['content']}<|im_end|>\n"
    out += "<|im_start|>assistant\n"
    return out


def extract_text(content) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
            elif isinstance(item, dict) and "text" in item:
                parts.append(item["text"])
        return " ".join(parts)
    return ""


# ─── Server ──────────────────────────────────────────────────────────────────

app = Server(title="Nova-1 Chat")


# ─── YaRN Context Stretch ────────────────────────────────────────────────────

@app.api(name="yarn", concurrency_limit=1, time_limit=60)
def yarn(ctx: str) -> str:
    """Apply YaRN RoPE scaling to extend the context window."""
    native_len = 2048
    if "2K" in ctx:
        target_len, factor = 2048, 1.0
    elif "4K" in ctx:
        target_len, factor = 4096, 2.0
    elif "8K" in ctx:
        target_len, factor = 8192, 4.0
    elif "16K" in ctx:
        target_len, factor = 16384, 8.0
    else:
        target_len, factor = 2048, 1.0

    if factor == 1.0:
        model.config.max_len = native_len
        if hasattr(model.config, "rope_scaling"):
            model.config.rope_scaling = None
        if hasattr(model, "_rope_cache"):
            model._rope_cache = None
        return "✅ Native 2K context active."

    model.config.max_len = target_len
    model.config.rope_scaling = {"type": "yarn", "factor": factor}
    if hasattr(model, "_rope_cache"):
        model._rope_cache = None

    k = target_len // 1024
    return f"✅ Stretched to {k}K (Factor {factor}.0)"


# ─── GPU Generation (separate from @app.api — do NOT stack @spaces.GPU+@app.api) ─

@spaces.GPU(duration=120)
def _generate_gpu(text, max_tokens, temperature, top_p, im_end_id, eos_id):
    """Runs on allocated ZeroGPU. Generator yields text chunks."""
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=model.config.max_len,
        add_special_tokens=False,
    )
    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
        timeout=60.0,
    )

    generation_kwargs = dict(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_k=50,
        top_p=top_p,
        do_sample=True,
        pad_token_id=eos_id,
        eos_token_id=[im_end_id, eos_id],
        repetition_penalty=1.15,
        streamer=streamer,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for chunk in streamer:
        yield chunk

    thread.join()


# ─── Generate Endpoint (streaming via SSE) ────────────────────────────────────

@app.api(name="generate", concurrency_limit=1, time_limit=180, stream_every=0.2)
def generate(
    message: str,
    history: list,
    system_prompt: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    mode: str,
):
    """Streaming chat generation. Yields text chunks via SSE."""
    try:
        if not message or not isinstance(message, str) or not message.strip():
            yield "⚠️ Please enter a message."
            return
        if len(message) > 8000:
            yield "⚠️ Message too long (8000 char max)."
            return

        clean_history = []
        for h in (history or []):
            if isinstance(h, dict):
                clean_history.append({
                    "role": h["role"],
                    "content": extract_text(h.get("content", "")),
                })

        max_hist_len = getattr(model.config, "max_len", 2048)
        if len(clean_history) > max_hist_len // 4:
            clean_history = clean_history[-(max_hist_len // 4):]

        if mode == "Chat":
            text = format_chatml(clean_history, system_prompt)
        else:
            text = ""
            if system_prompt and system_prompt.strip():
                text += system_prompt.strip() + "\n\n"
            for m in clean_history:
                text += m["content"] + "\n"
            text += message.strip()

        temperature = max(0.01, min(2.0, float(temperature)))
        top_p = max(0.01, min(1.0, float(top_p)))
        max_tokens = max(1, min(2048, int(max_tokens)))

        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        eos_id = tokenizer.convert_tokens_to_ids("</s>")
        if im_end_id == tokenizer.unk_token_id:
            im_end_id = 50258
        if eos_id == tokenizer.unk_token_id:
            eos_id = 50256

        yield from _generate_gpu(text, max_tokens, temperature, top_p, im_end_id, eos_id)

    except torch.cuda.OutOfMemoryError:
        torch.cuda.empty_cache()
        yield "⚠️ GPU OOM — try shorter input or fewer tokens."
    except Exception as e:
        print(f"❌ Generation error: {e}\n{traceback.format_exc()}")
        yield f"⚠️ Error: {e}"


# ─── Serve Custom Frontend ────────────────────────────────────────────────────

@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()


# ─── Launch ───────────────────────────────────────────────────────────────────

demo = app  # HF Spaces runtime expects variable named `demo`

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )