import os
import json
from typing import Any, Dict, List, Optional, Union

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM


MODEL_ID = os.getenv("MODEL_ID", "oleh13/ord-retro-qwen25-15b-merged-fp16")

hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)

app = FastAPI(
    title="ORD Retrosynthesis Model API",
    version="1.0.0",
)


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


class GenerateRequest(BaseModel):
    prompt: Optional[str] = None
    messages: Optional[List[ChatMessage]] = None

    max_new_tokens: int = Field(default=1200, ge=1, le=4096)
    temperature: float = Field(default=0.1, ge=0.0, le=2.0)
    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
    repetition_penalty: float = Field(default=1.05, ge=0.5, le=2.0)
    do_sample: bool = True

    return_json_only: bool = True


class GenerateResponse(BaseModel):
    text: str
    parsed_json: Optional[Union[Dict[str, Any], List[Any]]] = None
    raw_output: str
    model_id: str


def extract_json_object(text: str) -> str:
    text = text.strip()

    first_obj = text.find("{")
    last_obj = text.rfind("}")

    first_arr = text.find("[")
    last_arr = text.rfind("]")

    obj_valid = first_obj != -1 and last_obj != -1 and last_obj > first_obj
    arr_valid = first_arr != -1 and last_arr != -1 and last_arr > first_arr

    if obj_valid and arr_valid:
        if first_obj < first_arr:
            return text[first_obj:last_obj + 1]
        return text[first_arr:last_arr + 1]

    if obj_valid:
        return text[first_obj:last_obj + 1]

    if arr_valid:
        return text[first_arr:last_arr + 1]

    raise ValueError("No JSON object or array found in model output.")


print(f"Loading tokenizer: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


print(f"Loading model: {MODEL_ID}")

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
    device_map = "auto"
else:
    torch_dtype = torch.float32
    device_map = "cpu"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch_dtype,
    device_map=device_map,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

model.eval()

print("Model loaded.")
print("CUDA:", torch.cuda.is_available())


@app.get("/")
def root():
    return {
        "status": "ok",
        "model_id": MODEL_ID,
        "cuda": torch.cuda.is_available(),
    }


@app.get("/health")
def health():
    return {
        "status": "ok",
        "model_id": MODEL_ID,
        "cuda": torch.cuda.is_available(),
    }


@app.post("/chat/completions", response_model=GenerateResponse)
def generate(req: GenerateRequest):
    if req.messages and req.prompt:
        raise HTTPException(
            status_code=400,
            detail="Send either 'prompt' or 'messages', not both.",
        )

    if not req.messages and not req.prompt:
        raise HTTPException(
            status_code=400,
            detail="Send either 'prompt' or 'messages'.",
        )

    if req.messages:
        messages = [m.model_dump() for m in req.messages]

        prompt_text = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    else:
        prompt_text = req.prompt

    inputs = tokenizer(
        [prompt_text],
        return_tensors="pt",
    )

    if torch.cuda.is_available():
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=req.max_new_tokens,
            temperature=req.temperature,
            top_p=req.top_p,
            repetition_penalty=req.repetition_penalty,
            do_sample=req.do_sample,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )

    raw_output = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True,
    ).strip()

    parsed_json = None
    final_text = raw_output

    if req.return_json_only:
        try:
            json_text = extract_json_object(raw_output)
            parsed_json = json.loads(json_text)
            final_text = json.dumps(parsed_json, ensure_ascii=False, indent=2)
        except Exception as exc:
            raise HTTPException(
                status_code=422,
                detail={
                    "message": "Model did not return valid JSON.",
                    "error": str(exc),
                    "raw_output": raw_output,
                },
            )

    return GenerateResponse(
        text=final_text,
        parsed_json=parsed_json,
        raw_output=raw_output,
        model_id=MODEL_ID,
    )