# Cache bust commit: 20260720_v5
import sys
import traceback
from fastapi import FastAPI
from pydantic import BaseModel
import torch
import uvicorn
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

print("=== STARTING LEGALPARAM FASTAPI DOCKER SERVICE (FP16 v5) ===")
print("CUDA Available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("Device Name:", torch.cuda.get_device_name(0))

load_error = None
tokenizer = None
model = None

try:
    model_id = "bharatgenai/LegalParam"
    print("Loading config for", model_id)
    config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
    config.rope_scaling = {"type": "linear", "factor": 1.0}

    print("Loading tokenizer with trust_remote_code=False...")
    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=False)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    print("Loading model onto GPU with float16...")
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        config=config,
        trust_remote_code=True,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
    )
    model.eval()
    print("✅ Model loaded successfully onto device:", model.device)
except Exception as e:
    load_error = str(e)
    print("❌ ERROR LOADING MODEL:", load_error, file=sys.stderr)
    traceback.print_exc(file=sys.stderr)

app = FastAPI(title="LegalParam IBC API")


class SummaryRequest(BaseModel):
    text: str


@app.get("/")
def root():
    return {
        "status": "online" if model else "error",
        "model": "bharatgenai/LegalParam",
        "cuda": torch.cuda.is_available(),
        "load_error": load_error,
    }


@app.post("/summarize")
def summarize(req: SummaryRequest):
    if load_error or model is None:
        return {"error": f"Model not loaded: {load_error}"}

    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        prompt = (
            "You are a legal assistant. Summarize the following Indian Insolvency Law case into a concise legal digest with:\n"
            "## Facts\n## Issue & Holding\n## Key Principle\n\n"
            f"{req.text}\n\n"
            "Summary:\n"
        )

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            max_length=900,
            truncation=True,
        ).to(model.device)

        input_len = inputs.input_ids.shape[1]
        print(f"Input tokens: {input_len}", flush=True)

        eos_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 3
        pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else eos_id

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=500,
                do_sample=False,        # greedy — deterministic, no temp/top_p loops
                eos_token_id=eos_id,
                pad_token_id=pad_id,
                use_cache=True,         # KV cache ON — fast token generation
                repetition_penalty=1.1, # prevent repetition loops
            )

        print(f"Output tokens: {outputs.shape[1]}, new tokens: {outputs.shape[1] - input_len}", flush=True)

        new_tokens = outputs[0][input_len:]
        summary = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

        # Fallback: split on "Summary:" marker
        if not summary:
            raw_decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
            if "Summary:" in raw_decoded:
                summary = raw_decoded.split("Summary:")[-1].strip()
            else:
                summary = raw_decoded.strip()

        print(f"Summary length: {len(summary)} chars", flush=True)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        return {"summary": summary}

    except Exception as gen_err:
        err_msg = f"{type(gen_err).__name__}: {str(gen_err)}"
        print("❌ GENERATION ERROR:", err_msg, file=sys.stderr)
        return {"error": err_msg, "traceback": traceback.format_exc()}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)