import os
import re
import json
import torch
import torch.nn.functional as F
from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from pyvi import ViTokenizer
from PIL import Image
from torchvision import transforms
import timm
import requests
from flask_cors import CORS

# =====================
# CONFIG
# =====================
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"  # tránh vượt storage limit
TEXT_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Text_Image"
IMAGE_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Image"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
THRESHOLD = 0.65

# =====================
# LOAD TEXT MODEL
# =====================
print("🔹 Downloading text model...")

tokenizer = AutoTokenizer.from_pretrained(
    TEXT_MODEL_REPO,
    subfolder="text_model",
    use_fast=False
)

text_model = AutoModelForSequenceClassification.from_pretrained(
    TEXT_MODEL_REPO,
    subfolder="text_model/checkpoint-3390"
)
text_model.to(DEVICE).eval()

# load label2id.json
label_url = f"https://huggingface.co/{TEXT_MODEL_REPO}/resolve/main/text_model/label2id.json"
label2id = requests.get(label_url).json()
id2label_text = {i: l for l, i in label2id.items()}

# =====================
# LOAD IMAGE MODEL
# =====================
print("🔹 Downloading image model...")

class_names = ["an_toan", "bao_luc", "khieu_dam_doi_truy", "nhay_cam_chinh_tri"]

def build_model(num_classes=4):
    return timm.create_model("efficientnet_b3", pretrained=False, num_classes=num_classes)

image_model = build_model()

# tải model tạm trong /tmp để không chiếm storage
image_model_path = "/tmp/efficientnet_b3.pth"
if not os.path.exists(image_model_path):
    url = f"https://huggingface.co/{IMAGE_MODEL_REPO}/resolve/main/image_model/efficientnet_b3.pth"
    torch.hub.download_url_to_file(url, image_model_path)

image_model.load_state_dict(torch.load(image_model_path, map_location=DEVICE))
image_model.to(DEVICE).eval()

# chuẩn hóa ảnh
val_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# =====================
# UTILS
# =====================
def seg_pyvi(text: str) -> str:
    try:
        seg = ViTokenizer.tokenize(text)
        seg = seg.replace(" ", "_")
    except Exception:
        seg = text
    return seg

def split_sentences(text: str):
    sents = re.split(r'(?<=[.!?])\s+|\n+', text.strip())
    return [s for s in sents if s.strip()]

def predict_text(text: str):
    sentences = split_sentences(text)
    results = []
    for sent in sentences:
        seg = seg_pyvi(sent)
        inputs = tokenizer(seg, truncation=True, padding="max_length",
                           max_length=128, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            logits = text_model(**inputs).logits
            probs = F.softmax(logits, dim=-1).cpu().numpy()[0]
        pred_id = probs.argmax()
        label = id2label_text[int(pred_id)]
        prob = float(probs[pred_id])
        if label != "an_toan" and prob >= THRESHOLD:
            results.append({"sentence": sent, "label": label, "confidence": prob})
        else:
            results.append({
                "sentence": sent,
                "label": "an_toan",
                "confidence": float(probs[label2id["an_toan"]])
            })
    return results

def predict_image(pil_image: Image.Image):
    img = val_transforms(pil_image).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        outputs = image_model(img)
        probs = F.softmax(outputs, dim=1)[0].cpu().numpy()
    pred_id = probs.argmax()
    label = class_names[pred_id]
    prob = float(probs[pred_id])
    if label != "an_toan" and prob >= THRESHOLD:
        return {"label": label, "confidence": prob}
    else:
        return {
            "label": "an_toan",
            "confidence": float(probs[class_names.index("an_toan")])
        }

# =====================
# FLASK APP
# =====================
app = Flask(__name__)
CORS(app)

@app.route("/")
def home():
    return jsonify({"message": "✅ AI moderation API is running!"})

@app.route("/analyze", methods=["POST"])
def analyze():
    result = {"text_result": [], "image_result": []}

    # 🧠 PHÂN TÍCH TEXT
    if "content" in request.form:
        text_input = request.form["content"]
        result["text_result"] = predict_text(text_input)

    # 🧠 PHÂN TÍCH NHIỀU ẢNH
    if "image" in request.files:
        image_files = request.files.getlist("image")
        for image_file in image_files:
            image = Image.open(image_file.stream).convert("RGB")
            image_result = predict_image(image)
            result["image_result"].append(image_result)

    return jsonify(result)

# =====================
# RUN APP
# =====================
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))  # Hugging Face Space truyền PORT vào
    app.run(host="0.0.0.0", port=port)