from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer
from contextlib import asynccontextmanager
from minio import Minio
from minio.error import S3Error
import os
from typing import List

from bert_model import BERTClassifier
from utils import get_data_from_yaml, split_chunk
from preprocess import preprocess_text

config = get_data_from_yaml("/src/config.yaml")
DEVICE = config.get("device")
CLASSES = config.get("classes")
MINIO_SERVER = config.get("minio")["server"]
MINIO_DATA_LABELED = config.get("minio")["data_labeled"]
MINIO_MODEL_TRAINED = config.get("minio")["model_trained"]
VNCORENLP_DIR = config.get("vncorenlp")["save_dir"]
PHOBERTBASE_DIR = config.get("phobert_base")["save_dir"]
MAX_TOKEN_LENGTH = config.get("phobert_base")["max_token_length"]
MODEL_CHECKPOINT = config.get("model_checkpoint")
CHUNK_SIZE = config.get("chunk_size")
INFER_LENGTH = config.get("limit_infer_length")

minio_client = Minio(
    endpoint=MINIO_SERVER,
    access_key=os.getenv("MINIO_ROOT_USER"),
    secret_key=os.getenv("MINIO_ROOT_PASSWORD"),
    secure=False
)

def download_latest_model():
    # List objects in the bucket
    objects = minio_client.list_objects(MINIO_MODEL_TRAINED)
    latest_obj = None
    latest_time = None
    
    for obj in objects:
        if "best" in obj.object_name:
            if latest_time is None or obj.last_modified > latest_time:
                latest_time = obj.last_modified
                latest_obj = obj

    if latest_obj is not None:
        try:
            minio_client.fget_object(MINIO_MODEL_TRAINED, latest_obj.object_name, MODEL_CHECKPOINT)
        except S3Error as exc:
            print(f"Error occurred: {exc}")
        return latest_obj.object_name
    else:
        raise Exception("No *best* models found in the bucket")



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()
        print(f"Model updated: {update_model}")
    except Exception as e:
        print(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


