from fastapi import FastAPI, Request
from pydantic import BaseModel
import torch
import pickle
import gluonnlp as nlp
import numpy as np
import os
import sys
import collections
import logging # 로깅 모듈 임포트

from transformers import AutoTokenizer, BertModel
from torch.utils.data import Dataset, DataLoader
from huggingface_hub import hf_hub_download

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class BERTClassifier(torch.nn.Module):
    def __init__(self,
                 bert,
                 hidden_size = 768,
                 num_classes=5, # 분류할 클래스 수 (category 딕셔너리 크기와 일치)
                 dr_rate=None,
                 params=None):
        super(BERTClassifier, self).__init__()
        self.bert = bert
        self.dr_rate = dr_rate

        self.classifier = torch.nn.Linear(hidden_size , num_classes)
        if dr_rate:
            self.dropout = torch.nn.Dropout(p=dr_rate)

    def gen_attention_mask(self, token_ids, valid_length):
        attention_mask = torch.zeros_like(token_ids)
        for i, v in enumerate(valid_length):
            attention_mask[i][:v] = 1
        return attention_mask.float()

    def forward(self, token_ids, valid_length, segment_ids):
        attention_mask = self.gen_attention_mask(token_ids, valid_length)

        _, pooler = self.bert(input_ids=token_ids, token_type_ids=segment_ids.long(), attention_mask=attention_mask.float().to(token_ids.device), return_dict=False)

        if self.dr_rate:
            out = self.dropout(pooler)
        else:
            out = pooler
        return self.classifier(out)

# --- 2. BERTDataset 클래스 정의 ---
class BERTDataset(Dataset):
    def __init__(self, dataset, sent_idx, label_idx, bert_tokenizer, vocab, max_len, pad, pair):
        transform = nlp.data.BERTSentenceTransform(
            bert_tokenizer, max_seq_length=max_len, vocab=vocab, pad=pad, pair=pair
        )
        self.sentences = [transform([i[sent_idx]]) for i in dataset]
        self.labels = [np.int32(i[label_idx]) for i in dataset]

    def __getitem__(self, i):
        return (self.sentences[i] + (self.labels[i],))

    def __len__(self):
        return len(self.labels)

app = FastAPI()
device = torch.device("cpu") # Hugging Face Spaces의 무료 티어는 주로 CPU를 사용합니다.

# ✅ category 로드
try:
    with open("category.pkl", "rb") as f:
        category = pickle.load(f)
    logger.info("category.pkl 로드 성공.")
except FileNotFoundError:
    logger.error("Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
    sys.exit(1)

# ✅ vocab 로드
try:
    with open("vocab.pkl", "rb") as f:
        vocab = pickle.load(f)
    logger.info("vocab.pkl 로드 성공.")
except FileNotFoundError:
    logger.error("Error: vocab.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
    sys.exit(1)

# ✅ 토크나이저 로드
tokenizer = AutoTokenizer.from_pretrained('skt/kobert-base-v1')
logger.info("토크나이저 로드 성공.")

# ✅ 모델 로드 (Hugging Face Hub에서 다운로드)
try:
    HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
    HF_MODEL_FILENAME = "textClassifierModel.pt"
    
    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
    logger.info(f"모델 파일이 '{model_path}'에 성공적으로 다운로드되었습니다.")

    bert_base_model = BertModel.from_pretrained('skt/kobert-base-v1')
    model = BERTClassifier(
        bert_base_model,
        dr_rate=0.5, # 학습 시 사용된 dr_rate 값으로 변경하세요.
        num_classes=len(category)
    )

    loaded_state_dict = torch.load(model_path, map_location=device)
    
    new_state_dict = collections.OrderedDict()
    for k, v in loaded_state_dict.items():
        name = k
        if name.startswith('module.'):
            name = name[7:]
        new_state_dict[name] = v
    
    model.load_state_dict(new_state_dict, strict=False)
    model.to(device)
    model.eval()
    logger.info("모델 로드 성공.")

except Exception as e:
    logger.error(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
    sys.exit(1)


# ✅ 데이터셋 생성에 필요한 파라미터
max_len = 64
batch_size = 32

# ✅ 예측 함수
def predict(predict_sentence):
    data = [predict_sentence, '0']
    dataset_another = [data]
    another_test = BERTDataset(dataset_another, 0, 1, tokenizer.tokenize, vocab, max_len, True, False)
    test_dataLoader = DataLoader(another_test, batch_size=batch_size, num_workers=0)

    model.eval()

    with torch.no_grad():
        for batch_id, (token_ids, valid_length, segment_ids, label) in enumerate(test_dataLoader):
            token_ids = token_ids.long().to(device)
            segment_ids = segment_ids.long().to(device)
            
            out = model(token_ids, valid_length, segment_ids)
            
            logits = out # 모델의 직접 출력은 로짓입니다.
            probs = torch.nn.functional.softmax(logits, dim=1) # 확률 계산

            predicted_category_index = torch.argmax(probs, dim=1).item() # 예측 인덱스
            predicted_category_name = list(category.keys())[predicted_category_index] # 예측 카테고리 이름
            
            # --- 예측 상세 로깅 ---
            logger.info(f"Input Text: '{predict_sentence}'")
            logger.info(f"Raw Logits: {logits.tolist()}")
            logger.info(f"Probabilities: {probs.tolist()}")
            logger.info(f"Predicted Index: {predicted_category_index}")
            logger.info(f"Predicted Label: '{predicted_category_name}'")
            # --- 예측 상세 로깅 끝 ---

            return predicted_category_name

# ✅ 엔드포인트 정의
class InputText(BaseModel):
    text: str

@app.get("/")
def root():
    return {"message": "Text Classification API (KoBERT)"}

@app.post("/predict")
async def predict_route(item: InputText):
    result = predict(item.text)
    return {"text": item.text, "classification": result}