from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer
from contextlib import asynccontextmanager
from typing import List

from bert_model import BERTClassifier
from utils import get_data_from_yaml, split_chunk
from minio_client import download_latest_model
from constants import PHOBERTBASE_DIR, MAX_TOKEN_LENGTH, DEVICE, CLASSES, MODEL_CHECKPOINT, INFER_LENGTH, CHUNK_SIZE
from logger import logger
from preprocess import preprocess_text

tokenizer = AutoTokenizer.from_pretrained(PHOBERTBASE_DIR, local_files_only=True, use_fast=False)

def infer(text, model, tokenizer, class_names, max_len=MAX_TOKEN_LENGTH+2):
    encoded_review = tokenizer.encode_plus(
        text,
        max_length=max_len,
        truncation=True,
        add_special_tokens=True,
        padding='max_length',
        return_attention_mask=True,
        return_token_type_ids=False,
        return_tensors='pt',
    )

    input_ids = encoded_review['input_ids'].to(DEVICE)
    attention_mask = encoded_review['attention_mask'].to(DEVICE)

    output = model(input_ids, attention_mask)
    conf, y_pred = torch.max(output, dim=1)

    return conf, class_names[y_pred]


model = BERTClassifier(model_bert=PHOBERTBASE_DIR, n_classes=len(CLASSES))
model.to(DEVICE)

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model
    try:
        update_model = download_latest_model()
        model.load_state_dict(torch.load(MODEL_CHECKPOINT))
        model.eval()
        logger.info(f"Model updated: {update_model}")
    except Exception as e:
        logger.error(f"An error occurred: {e}")
    yield
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()

class ParagraphRequest(BaseModel):
    paragraph: str

class ChunkLabelResponse(BaseModel):
    chunk: str
    label: str
    confidence: float

app = FastAPI(lifespan=lifespan)

@app.post("/text-classify", response_model=List[ChunkLabelResponse])
async def process_paragraph(request: ParagraphRequest):
    if len(request.paragraph) > INFER_LENGTH:
        raise HTTPException(status_code=400, detail=f"Max length: {INFER_LENGTH}")
    chunks = split_chunk(request.paragraph, max_words=CHUNK_SIZE)
    response = []
    for chunk in chunks:
        text = preprocess_text(chunk)
        confidence, processed_label = infer(text, model, tokenizer, CLASSES)
        response.append(ChunkLabelResponse(chunk=chunk, label=processed_label, confidence=confidence))
    return response


