from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np
import os

app = FastAPI(title="PGOS Emotion Agent API", version="1.0.0")

# ── 설정 ───────────────────────────────────────────
HF_USERNAME = "nuguri01"
HF_TOKEN    = os.environ.get("HF_TOKEN", "")  # Spaces Secret으로 주입

MAJOR_CLASSES = ["분노계", "슬픔계", "불안계", "상처계", "당황계", "기쁨계"]

MINOR_STRUCTURE = {
    "분노계": {"labels": ["분노", "좌절", "짜증"],        "repo": "pgos-emotion-stage2-anger"},
    "슬픔계": {"labels": ["슬픔", "우울", "후회"],        "repo": "pgos-emotion-stage2-sadness"},
    "불안계": {"labels": ["불안", "걱정"],                "repo": "pgos-emotion-stage2-anxiety"},
    "상처계": {"labels": ["상처", "억울", "질투"],        "repo": "pgos-emotion-stage2-hurt"},
    "당황계": {"labels": ["수치심", "죄책감", "당황"],    "repo": "pgos-emotion-stage2-embarrass"},
    "기쁨계": {"labels": ["기쁨", "안도"],                "repo": "pgos-emotion-stage2-joy"},
}

PGOS_CLASSES = [
    "분노", "좌절", "짜증", "슬픔", "우울", "후회",
    "불안", "걱정", "상처", "억울", "질투",
    "수치심", "죄책감", "당황", "기쁨", "안도"
]
label2id = {l: i for i, l in enumerate(PGOS_CLASSES)}

device = "cuda" if torch.cuda.is_available() else "cpu"

# ── 모델 로드 (서버 시작 시 한 번만) ──────────────
print("📥 모델 로딩 중...")

tokenizer = AutoTokenizer.from_pretrained(
    f"{HF_USERNAME}/pgos-emotion-stage1-major",
    token=HF_TOKEN
)

model_major = AutoModelForSequenceClassification.from_pretrained(
    f"{HF_USERNAME}/pgos-emotion-stage1-major",
    token=HF_TOKEN
).to(device)
model_major.eval()

minor_models = {}
for major_label, config in MINOR_STRUCTURE.items():
    model_sub = AutoModelForSequenceClassification.from_pretrained(
        f"{HF_USERNAME}/{config['repo']}",
        token=HF_TOKEN,
        num_labels=len(config["labels"])
    ).to(device)
    model_sub.eval()
    minor_models[major_label] = model_sub

print("✅ 전체 모델 로드 완료")


# ── 요청/응답 스키마 ───────────────────────────────
class EmotionRequest(BaseModel):
    text: str
    top_k: int = 3  # 반환할 감정 수 (기본 3개)

class EmotionResult(BaseModel):
    emotion: str
    score: float
    rank: int

class EmotionResponse(BaseModel):
    text: str
    primary_emotion: str
    primary_score: float
    top_emotions: list[EmotionResult]
    pgos_message: str  # 내담자에게 돌려줄 멘트


# ── 감정 분류 함수 ─────────────────────────────────
def predict_emotions(text: str, top_k: int = 3):
    # 1단계: 대분류
    enc = tokenizer(text, max_length=128, padding=True,
                    truncation=True, return_tensors="pt").to(device)
    with torch.no_grad():
        major_probs = torch.softmax(
            model_major(**enc).logits, dim=-1
        ).cpu().numpy()[0]

    # 2단계: 소분류 + 확률 결합
    final_probs = np.zeros(len(PGOS_CLASSES))

    for major_idx, major_label in enumerate(MAJOR_CLASSES):
        config       = MINOR_STRUCTURE[major_label]
        minor_labels = config["labels"]
        model_sub    = minor_models[major_label]

        with torch.no_grad():
            sub_probs = torch.softmax(
                model_sub(**enc).logits, dim=-1
            ).cpu().numpy()[0]

        for minor_idx, minor_label in enumerate(minor_labels):
            final_probs[label2id[minor_label]] = (
                major_probs[major_idx] * sub_probs[minor_idx]
            )

    # Top-K 추출
    top_k_ids = np.argsort(final_probs)[-top_k:][::-1]
    top_emotions = [
        EmotionResult(
            emotion=PGOS_CLASSES[idx],
            score=round(float(final_probs[idx]), 4),
            rank=rank + 1
        )
        for rank, idx in enumerate(top_k_ids)
    ]

    # PGOS 멘트 생성
    primary  = top_emotions[0].emotion
    subs     = [e.emotion for e in top_emotions[1:] if e.score > 0.05]
    sub_str  = f", {'과 '.join(subs)} 감정도 함께 느껴지는 것 같아요" if subs else ""
    message  = f"'{primary}' 감정이 주로 느껴지며{sub_str}. 지금 이 감정을 천천히 들여다볼까요?"

    return top_emotions, primary, float(final_probs[top_k_ids[0]]), message


# ── API 엔드포인트 ─────────────────────────────────
@app.get("/")
def root():
    return {
        "service": "PGOS Emotion Agent API",
        "version": "1.0.0",
        "description": "계층적 2단계 감정 분류기 (16클래스, Top-3 Hit Rate 0.870)",
        "endpoints": {
            "POST /predict": "감정 분석",
            "GET /health": "서버 상태 확인"
        }
    }

@app.get("/health")
def health():
    return {"status": "ok", "device": device, "models_loaded": len(minor_models) + 1}

@app.post("/predict", response_model=EmotionResponse)
def predict(req: EmotionRequest):
    if not req.text.strip():
        raise HTTPException(status_code=400, detail="텍스트를 입력해주세요.")
    if len(req.text) > 500:
        raise HTTPException(status_code=400, detail="텍스트는 500자 이하로 입력해주세요.")

    top_k = max(1, min(req.top_k, 5))  # 1~5 사이로 제한

    top_emotions, primary, primary_score, message = predict_emotions(req.text, top_k)

    return EmotionResponse(
        text=req.text,
        primary_emotion=primary,
        primary_score=primary_score,
        top_emotions=top_emotions,
        pgos_message=message
    )