import time
import joblib
import numpy as np
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification, RobertaConfig, RobertaForSequenceClassification
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

SENTIMENT_MODEL_ID = "vanhai123/phobert-vi-comment-4class"
HATE_MODEL_ID = "visolex/visobert-hsd"
MAX_LENGTH = 128

torch.set_num_threads(2)  # Tối ưu hóa khớp với cấu hình 2 vCPU basic free của HF Space

app = FastAPI(title="Vietnamese Hate Speech Detection API")

print("Loading models... (chỉ chạy 1 lần khi Space khởi động)")

# 1. Load Sentiment Model
sent_tok = AutoTokenizer.from_pretrained(SENTIMENT_MODEL_ID)
sent_model = AutoModelForSequenceClassification.from_pretrained(SENTIMENT_MODEL_ID).eval()
sent_labels = [sent_model.config.id2label[i] for i in range(sent_model.config.num_labels)]

# 2. Load Hate Model (Áp dụng bản vá map cấu hình và trọng số thủ công)
hate_tok = AutoTokenizer.from_pretrained(HATE_MODEL_ID)

hate_config = RobertaConfig(
    vocab_size=15004,
    hidden_size=768,
    num_hidden_layers=12,
    num_attention_heads=12,
    intermediate_size=3072,
    max_position_embeddings=514,
    type_vocab_size=2,
    pad_token_id=1,
    bos_token_id=0,
    eos_token_id=2,
    num_labels=3,
    id2label={0: "CLEAN", 1: "OFFENSIVE", 2: "HATE"},
    label2id={"CLEAN": 0, "OFFENSIVE": 1, "HATE": 2}
)

hate_model = RobertaForSequenceClassification(hate_config)

# Tải file trọng số gốc về Space ảo
weights_path = hf_hub_download(repo_id=HATE_MODEL_ID, filename="model.safetensors")
state_dict = load_file(weights_path)

# Đổi lại tiền tố biến từ encoder sang roberta
fixed_state_dict = {}
for key, value in state_dict.items():
    new_key = key
    if key.startswith("encoder."):
        new_key = key.replace("encoder.", "roberta.", 1)
    fixed_state_dict[new_key] = value

hate_model.load_state_dict(fixed_state_dict, strict=False)
hate_model.eval()
hate_labels = [hate_model.config.id2label[i] for i in range(hate_model.config.num_labels)]

# 3. Load Meta-Classifier từ file joblib (1.28 KB)
meta_bundle = joblib.load("meta_classifier.joblib")
meta_clf = meta_bundle["model"]

print("Models loaded successfully. Ready to serve.")


@torch.no_grad()
def get_probs(text, tok, model):
    inputs = tok(text, truncation=True, max_length=MAX_LENGTH, return_tensors="pt")
    logits = model(**inputs).logits
    probs = torch.softmax(logits, dim=-1).squeeze(0).numpy()
    return probs


class PredictRequest(BaseModel):
    text: str


class PredictResponse(BaseModel):
    final_label: str
    confidence: float
    sentiment: dict
    hate_speech: dict
    latency_ms: float


@app.get("/")
def health():
    return {"status": "ok"}


@app.post("/predict", response_model=PredictResponse)
def predict(req: PredictRequest):
    start = time.time()

    # Trích xuất phân phối xác suất từ 2 model lõi
    p_sent = get_probs(req.text, sent_tok, sent_model)
    p_hate = get_probs(req.text, hate_tok, hate_model)

    # Lấy label dự đoán của riêng từng model độc lập
    sent_label = sent_labels[int(np.argmax(p_sent))]
    hate_label = hate_labels[int(np.argmax(p_hate))]

    # Ghép chuỗi xác suất thành vector đặc trưng (7 chiều) truyền vào Meta Classifier
    features = np.concatenate([p_sent, p_hate]).reshape(1, -1)
    final_pred = meta_clf.predict(features)[0]
    final_proba = meta_clf.predict_proba(features)[0]
    confidence = float(np.max(final_proba))

    elapsed_ms = (time.time() - start) * 1000

    return {
        "final_label": final_pred,
        "confidence": round(confidence, 4),
        "sentiment": {
            "label": sent_label,
            "scores": {l: round(float(s), 4) for l, s in zip(sent_labels, p_sent)},
        },
        "hate_speech": {
            "label": hate_label,
            "scores": {l: round(float(s), 4) for l, s in zip(hate_labels, p_hate)},
        },
        "latency_ms": round(elapsed_ms, 1),
    }