import os
import sys
import time
import json
import asyncio
from typing import List, Optional, AsyncGenerator

import torch
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
from transformers import PreTrainedTokenizerFast

from model_def import PurcarModel, SEQ_LEN, ROPE_SCALE
from lora_layer import apply_lora, merge_lora_to_base

# Fix console encodings on Windows
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        sys.stderr.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

CHECKPOINT_PATH = os.environ.get('PURCAR_CHECKPOINT', r"C:\Users\kenyp\Documents\purcarbbc\checkpoint.pt")
BEST_PATH = os.environ.get('PURCAR_WEIGHTS', r"C:\Users\kenyp\Documents\purcarbbc\best.pt")

MODEL_PATHS = [
    CHECKPOINT_PATH,
    BEST_PATH,
    os.path.join(os.path.dirname(__file__), "weights", "checkpoint.pt"),
    os.path.join(os.path.dirname(__file__), "weights", "best.pt"),
]

TOKENIZER_PATHS = [
    os.environ.get('PURCAR_TOKENIZER', '/app/backend/en_tokenizer_hf'),
    r"C:\Users\kenyp\Documents\purcarbbc\en_tokenizer_hf",
    os.path.join(os.path.dirname(__file__), "en_tokenizer_hf"),
]

app = FastAPI(title="Purcar-1 Sovereign Engine", version="2.1.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

device = "cuda" if torch.cuda.is_available() else "cpu"
model: Optional[PurcarModel] = None
tokenizer: Optional[PreTrainedTokenizerFast] = None

pad_id = 0
unk_id = 1
bos_id = 2
eos_id = 3
context_id = 4
user_id = 5
assistant_id = 6

lora_merged_status = False
loaded_epoch = 0
loaded_step = 0
loaded_loss = 0.0


def find_first(paths):
    for p in paths:
        if os.path.exists(p):
            return p
    return None


def load_engine():
    global model, tokenizer, bos_id, eos_id, pad_id, user_id, assistant_id
    global lora_merged_status, loaded_epoch, loaded_step, loaded_loss

    print("=" * 65)
    print(f"[*] Initializing Purcar-1 Sovereign Model Engine on {device}")
    if torch.cuda.is_available():
        print(f"[*] GPU Device: {torch.cuda.get_device_name(0)}")
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.backends.cudnn.benchmark = True
    print(f"[*] Max Context Window: {SEQ_LEN} tokens (RoPE scale: {ROPE_SCALE:.2f}x)")
    print("=" * 65)

    tok_path = find_first(TOKENIZER_PATHS)
    if not tok_path:
        print("[!] Tokenizer not found in paths:", TOKENIZER_PATHS)
        return False
    print("[*] Loading tokenizer from", tok_path)
    tokenizer = PreTrainedTokenizerFast(tokenizer_file=os.path.join(tok_path,'tokenizer.json'),
        bos_token='<bos>',eos_token='<eos>',pad_token='<pad>',unk_token='<unk>',
        additional_special_tokens=['<context>','<user>','<assistant>'])
    bos_id = tokenizer.bos_token_id or 2
    eos_id = tokenizer.eos_token_id or 3
    pad_id = tokenizer.pad_token_id or 0
    user_id = tokenizer.convert_tokens_to_ids("<user>") or 5
    assistant_id = tokenizer.convert_tokens_to_ids("<assistant>") or 6
    print(f"[*] Special Tokens: bos={bos_id}, eos={eos_id}, pad={pad_id}, user={user_id}, assistant={assistant_id}")

    # 1. Instantiate FlashAttention 20-Layer Transformer with 8192 scaled RoPE
    print("[1/3] Instantiating PurcarModel (236M params, FlashAttention)...")
    model = PurcarModel(
        d_model=896,
        nhead=14,
        dim_feedforward=2560,
        num_layers=20,
        vocab_size=32000,
        dropout=0.0,
        scale=ROPE_SCALE,
        pad_id=pad_id,
    ).to(device)

    # The published weights are a LoRA-fused training checkpoint. Apply the
    # adapter structure before loading it; loading it as a plain base model
    # produces hundreds of missing/unexpected keys.
    ckpt_file = CHECKPOINT_PATH if os.path.exists(CHECKPOINT_PATH) else BEST_PATH
    if not ckpt_file or not os.path.exists(ckpt_file):
        raise RuntimeError('Model weights missing')
    if ckpt_file and os.path.exists(ckpt_file):
        print(f"[2/2] Loading trained weights from '{ckpt_file}'...")
        ck_chk = torch.load(ckpt_file, map_location=device, weights_only=True)
        state_chk = ck_chk["model"] if isinstance(ck_chk, dict) and "model" in ck_chk else ck_chk
        has_lora = any("lora_" in k for k in state_chk.keys())

        if has_lora:
            apply_lora(model, r=16, lora_alpha=32.0, lora_dropout=0.0)
            missing, unexpected = model.load_state_dict(state_chk, strict=False)
            if len(missing) > 20 or len(unexpected) > 20:
                raise RuntimeError(f'Model checkpoint mismatch ({len(missing)} missing, {len(unexpected)} unexpected)')
            loaded_epoch = ck_chk.get("epoch", 15) if isinstance(ck_chk, dict) else 15
            loaded_step = ck_chk.get("step", 1516000) if isinstance(ck_chk, dict) else 1516000
            loaded_loss = ck_chk.get("loss", 2.84) if isinstance(ck_chk, dict) else 2.84
            print(f"[CHECKPOINT INFO] Epoch: {loaded_epoch}, Step: {loaded_step}, Loss: {loaded_loss:.4f}" if isinstance(loaded_loss, (int, float)) else f"[CHECKPOINT INFO] Step: {loaded_step}")
            
            # Fuse LoRA mathematically directly into base weights
            merge_lora_to_base(model)
            lora_merged_status = True
            print("[HARDWARE ACCEL] LoRA merged into Tensor Core weights for maximum inference speed!")
        else:
            print("[LoRA] Checkpoint has full base weights without separate LoRA adapters.")

    model.eval()
    print("[OK] Purcar-1 Sovereign Engine is ready for 8192-token generation!\n")
    return True


@app.on_event("startup")
def startup_event():
    global model
    try:
        load_engine()
    except Exception as e:
        model = None
        print("[!] Startup error:", e)


class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: str = "purcar"
    messages: Optional[List[ChatMessage]] = None
    prompt: Optional[str] = None
    temperature: float = Field(0.4, ge=0.0, le=2.0)
    top_k: int = Field(20, ge=1, le=200)
    top_p: float = Field(0.9, ge=0.1, le=1.0)
    repetition_penalty: float = Field(1.15, ge=1.0, le=2.0)
    max_tokens: int = Field(2048, ge=1, le=8192)
    stream: bool = True


def build_input_ids(raw_prompt_or_messages) -> List[int]:
    """
    Constructs the canonical input token sequence:
    [bos_id] + tokenizer.encode(formatted_prompt, add_special_tokens=False)
    """
    if isinstance(raw_prompt_or_messages, str):
        prompt_text = raw_prompt_or_messages.strip()
    else:
        # Assemble from messages if not already formatted
        system_text = ""
        context_text = ""
        user_text = ""
        for m in raw_prompt_or_messages:
            if m.role == "system":
                system_text = m.content.strip()
            elif m.role == "context":
                context_text = m.content.strip()
            elif m.role == "user":
                user_text = m.content.strip()

        # If user text already has canonical tags, use directly
        if "<user>" in user_text:
            prompt_text = user_text
        else:
            ctx = context_text or system_text
            if ctx:
                prompt_text = f"<context> {ctx} Task: Generation <user> {user_text} <assistant>"
            else:
                prompt_text = f"<context> Task: Generation <user> {user_text} <assistant>"

    tokens = tokenizer.encode(prompt_text, add_special_tokens=False)
    return [bos_id] + tokens


@torch.inference_mode()
def sample_next_token(logits, temp=0.4, topk=20, topp=0.9, seen=None, penalty=1.15) -> int:
    """
    Applies repetition penalty, temperature scaling, top-k truncation, and top-p (nucleus) sampling.
    """
    if seen:
        for t in set(seen):
            if logits[0, t] < 0:
                logits[0, t] = logits[0, t] * penalty
            else:
                logits[0, t] = logits[0, t] / penalty
        if len(seen) >= 2 and seen[-1] == seen[-2]:
            logits[0, seen[-1]] -= 5.0

    logits = logits / max(temp, 1e-5)

    # Top-K
    top_values, top_indices = torch.topk(logits, min(topk, logits.size(-1)), dim=-1)

    # Top-P (Nucleus filtering) on top_values
    if topp < 1.0:
        probs = torch.softmax(top_values, dim=-1)
        cum_probs = torch.cumsum(probs, dim=-1)
        mask = cum_probs - probs > topp
        top_values[mask] = -float("inf")

    probs = torch.softmax(top_values, dim=-1)
    idx = torch.multinomial(probs, 1)
    return top_indices.gather(-1, idx).item()


@torch.inference_mode()
def generate_full(prompt_input, temp=0.4, topk=20, topp=0.9, penalty=1.15, max_length=2048):
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    ids = build_input_ids(prompt_input)
    # Available tokens remaining in the 8192 context window
    budget = min(max_length, SEQ_LEN - len(ids) - 1)
    if budget <= 0:
        raise HTTPException(status_code=413, detail="Prompt exceeds the model context window")

    generated = []
    seen = []

    for _ in range(budget):
        x = torch.tensor([ids], dtype=torch.long, device=device)
        with torch.amp.autocast("cuda") if device == "cuda" else torch.no_grad():
            logits = model(x)[:, -1, :]

        next_tok = sample_next_token(logits, temp=temp, topk=topk, topp=topp, seen=seen, penalty=penalty)
        if next_tok == eos_id:
            break

        ids.append(next_tok)
        seen.append(next_tok)
        generated.append(next_tok)

    return tokenizer.decode(generated, skip_special_tokens=True).strip()


async def generate_stream(
    prompt_input, temp=0.4, topk=20, topp=0.9, penalty=1.15, max_length=2048
):
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    def sse(text):
        payload = {"choices": [{"delta": {"content": text}, "finish_reason": None}]}
        return "data: " + json.dumps(payload) + "\n\n"

    header = {
        "id": f"chatcmpl-{int(time.time()*1000)}",
        "object": "chat.completion.chunk",
        "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}],
    }
    yield "data: " + json.dumps(header) + "\n\n"

    ids = build_input_ids(prompt_input)
    budget = min(max_length, SEQ_LEN - len(ids) - 1)
    if budget <= 0:
        raise HTTPException(status_code=413, detail="Prompt exceeds the model context window")

    generated_ids = []
    seen = []
    prev_text = ""

    for _ in range(budget):
        x = torch.tensor([ids], dtype=torch.long, device=device)
        with torch.amp.autocast("cuda") if device == "cuda" else torch.no_grad():
            logits = model(x)[:, -1, :]

        next_tok = sample_next_token(logits, temp=temp, topk=topk, topp=topp, seen=seen, penalty=penalty)
        if next_tok == eos_id:
            break

        ids.append(next_tok)
        seen.append(next_tok)
        generated_ids.append(next_tok)

        cur_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
        new_text = cur_text[len(prev_text):]
        if new_text:
            yield sse(new_text)
            prev_text = cur_text
        await asyncio.sleep(0.002)

    yield "data: " + json.dumps({"choices": [{"delta": {}, "finish_reason": "stop"}], "usage": {"prompt_tokens": len(build_input_ids(prompt_input)), "completion_tokens": len(generated_ids)}}) + "\n\n"
    yield "data: [DONE]\n\n"


@app.get("/health")
def health():
    param_count = sum(p.numel() for p in model.parameters()) if model else 0
    return {
        "status": "healthy" if model is not None else "unavailable",
        "device": device,
        "cuda_available": torch.cuda.is_available(),
        "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU",
        "model_loaded": model is not None,
        "lora_merged": lora_merged_status,
        "epoch": loaded_epoch,
        "step": loaded_step,
        "loss": loaded_loss,
        "context_window": SEQ_LEN,
        "rope_scale": ROPE_SCALE,
        "vocab_size": 32000,
        "param_count": param_count,
        "bos_id": bos_id,
        "eos_id": eos_id,
        "user_id": user_id,
        "assistant_id": assistant_id,
    }


@app.get("/v1/models")
def list_models():
    return {
        "data": [
            {
                "id": "purcar",
                "name": "Purcar-1 (8192 Context)",
                "context_length": SEQ_LEN,
                "owned_by": "purcarbbc",
            }
        ]
    }


@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
    # Support direct prompt or messages array
    prompt_input = req.prompt if req.prompt else req.messages
    if not prompt_input:
        raise HTTPException(status_code=400, detail="Missing prompt or messages")
    if model is None or tokenizer is None:
        raise HTTPException(status_code=503, detail="Model not loaded")
    if len(build_input_ids(prompt_input)) >= SEQ_LEN - 1:
        raise HTTPException(status_code=413, detail="Prompt exceeds the model context window")

    if req.stream:
        return StreamingResponse(
            generate_stream(
                prompt_input,
                temp=req.temperature,
                topk=req.top_k,
                topp=req.top_p,
                penalty=req.repetition_penalty,
                max_length=req.max_tokens,
            ),
            media_type="text/event-stream",
        )

    result = generate_full(
        prompt_input,
        temp=req.temperature,
        topk=req.top_k,
        topp=req.top_p,
        penalty=req.repetition_penalty,
        max_length=req.max_tokens,
    )
    return {
        "usage": {"prompt_tokens": len(build_input_ids(prompt_input)), "completion_tokens": len(tokenizer.encode(result, add_special_tokens=False))},
        "choices": [{"message": {"role": "assistant", "content": result}, "finish_reason": "stop"}]
    }


class RefineRequest(BaseModel):
    prompt: str
    language: Optional[str] = None


@app.post('/v1/tokenize')
def count_tokens(req: RefineRequest):
    if tokenizer is None or model is None:
        raise HTTPException(status_code=503, detail='Model not loaded')
    return {'tokens': len(build_input_ids(req.prompt))}


refiner_model = None
refiner_tokenizer = None


@app.post("/v1/refine")
def refine_prompt(req: RefineRequest):
    """
    Neural Prompt Refiner & Intent Classifier powered by local Qwen.
    Corrects spelling, typos, and broken slang without hardcoded dictionaries.
    """
    global refiner_model, refiner_tokenizer
    prompt = req.prompt.strip()
    if not prompt:
        return {"refined_prompt": "", "intent": "general", "search_query": ""}

    try:
        if refiner_model is None:
            from transformers import AutoTokenizer, AutoModelForCausalLM
            ref_id = "Qwen/Qwen2.5-0.5B-Instruct"
            refiner_tokenizer = AutoTokenizer.from_pretrained(ref_id)
            refiner_model = AutoModelForCausalLM.from_pretrained(
                ref_id,
                torch_dtype=torch.float16 if device == 'cuda' else torch.float32,
            ).to(device)

        system_prompt = (
            "You are a prompt normalizer and intent classifier. "
            "Task: Correct spelling only, preserving the exact meaning and language. Preserve slang, quoted text, names, and profanity. Do not answer the request, moralize, or replace it with a refusal. If no correction is needed, return the original text. "
            "Output strictly JSON with keys: 'refined_prompt', 'intent' (one of 'code', 'reasoning', 'factual_search', 'greeting', 'general'), 'search_query' (clean search query or empty string)."
        )

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]
        text = refiner_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = refiner_tokenizer([text], return_tensors="pt").to(device)

        with torch.no_grad():
            outputs = refiner_model.generate(**inputs, max_new_tokens=70, do_sample=False)
        gen = refiner_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        usage = int(inputs.input_ids.shape[1] + outputs.shape[1] - inputs.input_ids.shape[1])

        if "{" in gen and "}" in gen:
            json_str = gen[gen.find("{"):gen.rfind("}")+1]
            data = json.loads(json_str)
            return {
                "refined_prompt": data.get("refined_prompt", prompt),
                "intent": data.get("intent", "general"),
                "search_query": data.get("search_query", ""),
                "neural": True
                ,"usage_tokens": usage
            }
        return {"refined_prompt": prompt, "intent": "general", "search_query": "", "neural": False, "usage_tokens": usage}
    except Exception as e:
        return {"refined_prompt": prompt, "intent": "general", "search_query": "", "neural": False, "error": str(e)}


if __name__ == "__main__":
    import uvicorn
    print(f"[*] Starting Purcar-1 Sovereign Server on http://0.0.0.0:8000")
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False)