from fastapi import FastAPI, Request
from transformers import BertModel, BertForSequenceClassification, AutoTokenizer
from huggingface_hub import hf_hub_download
import torch
import pickle
import os
import sys
import psutil

app = FastAPI()
device = torch.device("cpu")

# category.pkl 로드
try:
    with open("category.pkl", "rb") as f:
        category = pickle.load(f)
    print("✅ category.pkl 로드 성공.")
except FileNotFoundError:
    print("❌ Error: category.pkl 파일을 찾을 수 없습니다.")
    sys.exit(1)

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

class CustomClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # 정의했던 구조 그대로 복원해야 함
        self.bert = BertModel.from_pretrained("skt/kobert-base-v1")
        self.classifier = torch.nn.Linear(768, len(category))

    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
        outputs = self.bert(input_ids=input_ids,
                            attention_mask=attention_mask,
                            token_type_ids=token_type_ids)
        pooled_output = outputs[1]  # CLS 토큰
        return self.classifier(pooled_output)

HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
HF_MODEL_FILENAME = "textClassifierModel.pt"

# 메모리 측정 전
process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / (1024 * 1024)
print(f"📦 모델 다운로드 전 메모리 사용량: {mem_before:.2f} MB")

# 모델 로드
try:
    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
    print(f"✅ 모델 파일 다운로드 성공: {model_path}")

    state_dict = torch.load(model_path, map_location=device)
    model = BertForSequenceClassification.from_pretrained(
        "skt/kobert-base-v1",
        num_labels=len(category),
        state_dict=state_dict,
    )
    model.to(device)
    model.eval()
    print("✅ 모델 로드 및 준비 완료.")
except Exception as e:
    print(f"❌ Error: 모델 로드 중 오류 발생: {e}")
    sys.exit(1)


# 예측 API
@app.post("/predict")
async def predict_api(request: Request):
    data = await request.json()
    text = data.get("text")
    if not text:
        return {"error": "No text provided", "classification": "null"}

    encoded = tokenizer.encode_plus(
        text, max_length=64, padding='max_length', truncation=True, return_tensors='pt'
    )

    with torch.no_grad():
        outputs = model(**encoded)
        probs = torch.nn.functional.softmax(outputs.logits, dim=1)
        predicted = torch.argmax(probs, dim=1).item()

    label = list(category.keys())[predicted]
    return {"text": text, "classification": label}