import os
import re
import sqlite3
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from underthesea import text_normalize
from datetime import datetime

# --- Cấu hình model ---
MODEL_DIR = "5CD-AI/Vietnamese-Sentiment-visobert"
LABELS = {0: "negative", 1: "positive", 2: "neutral"}

# --- Bộ phân loại cảm xúc ---
class EmotionClassifier:
    def __init__(self, model_dir=MODEL_DIR, device=None):
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
            self.model = AutoModelForSequenceClassification.from_pretrained(model_dir)
        except Exception as e:
            raise RuntimeError(f"Không thể load model từ {model_dir}: {e}")
        self.model.to(self.device)
        self.model.eval()

    # --- Tiền xử lý cơ bản tiếng Việt ---
    def preprocess(self, text: str) -> str:
        if not isinstance(text, str):
            text = str(text)
        t = text.lower()
        # Thay thế các từ viết tắt, slang
        replacements = {
            r"\brat\b": "rất",
            r"\br\b": "rất",
            r"\bko\b": "không",
            r"\bk\b": "không",
            r"\bkhong\b": "không",
            r"\bkhog\b": "không",
            r"\bhok\b": "không",
            r"\bkg\b": "không",
            r"\bkh\b": "không",
            r"\bdo\b": "dở",
            r"\bbiet\b": "biết",
            r"\bbt\b": "biết",
            r"\bdc\b": "được",
            r"\bdk\b": "được",
            r"\bđc\b": "được",
            r"\bvs\b": "với",
            r"\bj\b": "gì",
            r"\bjz\b": "gì",
            r"\bgi\b": "gì",
            r"\bntn\b": "như thế nào",
            r"\bnt\b": "nhắn tin",
            r"\bmn\b": "mọi người",
            r"\bbth\b": "bình thường",
            r"\bbthg\b": "bình thường",
            r"\bthik\b": "thích",
            r"\bxin loi\b": "xin lỗi",
            r"\bxl\b": "xin lỗi",       
            r"\blm\b": "làm",
            r"\bcam on\b": "cảm ơn",
            r"\bcs\b": "có",
            r"\bco\b": "có",
            r"\bko the\b": "không thể",
            r"\btam\b": "tạm",
            r"\bcx\b": "cũng",
            r"\bcung\b": "cũng",
            r"\bcũm\b": "cũng",
            r"\bhn\b": "hôm nay",
            r"\bhom nay\b": "hôm nay",
            r"\bthui\b": "thôi",  
            r"\blun\b": "luôn",  
            r"\bng\b": "người",
            r"\bghek\b": "ghét"
        }
        for pat, repl in replacements.items():
            t = re.sub(pat, repl, t, flags=re.IGNORECASE)
        try:
            t = text_normalize(t)
        except Exception:
            pass
        t = re.sub(r'\s+', ' ', t).strip()
        return t[:50]

    def predict(self, text: str) -> dict:
        # --- Trả về dict: {label_id, sentiment, probs} ---
        t = self.preprocess(text)
        inputs = self.tokenizer(t, truncation=True, padding=True, return_tensors="pt", max_length=128)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits.cpu().numpy()[0]
            probs = torch.nn.functional.softmax(torch.tensor(logits), dim=-1).numpy()
            pred_id = int(np.argmax(logits))
        return {
            "label_id": pred_id,
            "sentiment": LABELS.get(pred_id, str(pred_id)),
            "probs": probs.tolist()
        }

# --- Lưu trữ lịch sử phân loại ---
class HistoryDB:
    def __init__(self, db_path=None):
        base = os.path.dirname(os.path.abspath(__file__))
        self.db_path = db_path or os.path.join(base, "history.db")
        self.conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self.create_table()

    def create_table(self):
        cur = self.conn.cursor()
        cur.execute("""
            CREATE TABLE IF NOT EXISTS history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                text TEXT NOT NULL,
                sentiment TEXT NOT NULL,
                timestamp DATETIME DEFAULT CURRENT_TIMESTAMP NOT NULL
            )
        """)
        self.conn.commit()

    def add(self, text: str, sentiment: str) -> int:
        cur = self.conn.cursor()
        cur.execute(
            "INSERT INTO history (text, sentiment, timestamp) VALUES (?, ?, datetime('now','localtime'))",
            (text, sentiment)
        )
        self.conn.commit()
        return cur.lastrowid

    def list_all(self, limit=200):
        cur = self.conn.cursor()
        cur.execute(
            "SELECT id, text, sentiment, timestamp FROM history ORDER BY id DESC LIMIT ?",
            (limit,)
        )
        return cur.fetchall()

# --- Self-test nhanh ---
if __name__ == "__main__":
    print("Testing app_backend.py...")
    try:
        clf = EmotionClassifier()
        print("Model loaded. Device:", clf.device)
        r = clf.predict("Mình rất vui hôm nay!")
        print("Sample predict:", r)
    except Exception as e:
        print("Model load/predict failed:", repr(e))
    try:
        db = HistoryDB()
        print("DB ready. Recent rows:", db.list_all(5))
        dbid = db.add("Mình rất vui hôm nay!", r["sentiment"])
        print("Inserted id:", dbid)
    except Exception as e:
        print("DB failed:", repr(e))