from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import os
import gc
import torch
import traceback

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

token = os.environ.get("HF_TOKEN")

class SummarizeRequest(BaseModel):
    text: str
    model_type: str

@app.get("/")
def read_root():
    return {"message": "Text Summarizer API Backend is Running!"}

@app.post("/api/summarize")
def summarize(req: SummarizeRequest):
    models = {
        "bart": "kUrrooYuki/bart-large-cnn-text-summarize",
        "t5": "kUrrooYuki/t5-base-text-summarize",
        "led": "kUrrooYuki/led-base-text-summarize"
    }
    
    if req.model_type not in models:
        raise HTTPException(status_code=400, detail="Model not recognized")

    try:
        model_repo = models[req.model_type]
        print(f"-> [1/3] Memuat {model_repo} ke RAM...", flush=True)
        
        tokenizer = AutoTokenizer.from_pretrained(model_repo, token=token)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_repo, token=token)
        
        text_input = req.text
        if req.model_type == "t5":
            text_input = "summarize: " + text_input
            
        print("-> [2/3] Starting the AI computation process (please wait)...", flush=True)
        inputs = tokenizer(text_input, return_tensors="pt", max_length=1024, truncation=True)

        in_ids = inputs["input_ids"]
        att_mask = inputs["attention_mask"]
        
        if not torch.is_tensor(in_ids):
            if isinstance(in_ids, list) and (len(in_ids) == 0 or not isinstance(in_ids[0], list)):
                in_ids = [in_ids]
            in_ids = torch.tensor(in_ids)
            
        if not torch.is_tensor(att_mask):
            if isinstance(att_mask, list) and (len(att_mask) == 0 or not isinstance(att_mask[0], list)):
                att_mask = [att_mask]
            att_mask = torch.tensor(att_mask)

        beams_config = 2 if req.model_type == "bart" else 1

        ngram_config = 3 if req.model_type == "bart" else 0 

        summary_ids = model.generate(
            input_ids=in_ids,               
            attention_mask=att_mask,     
            max_new_tokens=100,  
            min_length=20,
            num_beams=beams_config,
            no_repeat_ngram_size=ngram_config,
            early_stopping=True
        )
        
        print("-> [3/3] Computation complete! Translating tokens to text...", flush=True)
        summary_result = tokenizer.decode(summary_ids[0], skip_special_tokens=True)

        summary_result = summary_result.strip()
        
        if summary_result.lower().startswith("a's "):
            summary_result = summary_result[4:].strip()
        elif summary_result.lower().startswith("a "):
            summary_result = summary_result[2:].strip()
            
        if len(summary_result) > 0:
            summary_result = summary_result[0].upper() + summary_result[1:]
        
        del tokenizer
        del model
        gc.collect()
        print("-> Success. The RAM has been cleared.", flush=True)
        
        return {"status": "success", "summary": summary_result}
        
    except Exception as e:
        print("=== FULL ERROR TRACEBACK ===", flush=True)
        traceback.print_exc()
        print("============================", flush=True)
        raise HTTPException(status_code=500, detail=str(e))