"""
Chat trực tiếp qua terminal.
Tích hợp BPE Tokenizer + Cosine Similarity.
Chạy: python chat.py
"""
import torch
import torch.nn.functional as F
import pickle
import os
import re
import sys
import csv
from src.model import make_model

sys.stdout.reconfigure(encoding='utf-8')

VOCAB_PATH     = 'data/processed/vocab.pkl'
MODEL_PATH     = 'models/best/model.pt'
TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'
RAW_CSV_PATH   = 'data/raw/output.csv'

SIMILARITY_THRESHOLD = 0.92


def load_bpe_tokenizer():
    if os.path.exists(TOKENIZER_PATH):
        from tokenizers import Tokenizer
        return Tokenizer.from_file(TOKENIZER_PATH)
    return None


def tokenize_bpe(text, bpe_tok, word2idx):
    if bpe_tok is not None:
        return bpe_tok.encode(str(text).strip()).ids
    else:
        text = str(text).lower()
        text = re.sub(r'([.,!?()])', r' \1 ', text)
        text = re.sub(r'\s{2,}', ' ', text)
        return [word2idx.get(w, word2idx.get('<unk>', 3)) for w in text.strip().split()]


def build_qa_embeddings(model, device, bpe_tok, word2idx):
    qa_database = []
    if not os.path.exists(RAW_CSV_PATH):
        return qa_database, None

    with open(RAW_CSV_PATH, 'r', encoding='utf-8-sig') as f:
        reader = csv.reader(f)
        next(reader, None)
        for row in reader:
            if len(row) >= 3 and row[1].strip() and row[2].strip():
                qa_database.append({
                    "topic": row[0].strip(),
                    "question": row[1].strip(),
                    "answer": row[2].strip(),
                })

    if not qa_database:
        return qa_database, None

    sos_id = word2idx.get('<sos>', 1)
    all_embs = []
    for item in qa_database:
        q_ids = tokenize_bpe(item['question'], bpe_tok, word2idx)
        x = torch.tensor([[sos_id] + q_ids]).to(device)
        emb = model.get_embedding(x)
        all_embs.append(emb)

    qa_embeddings = F.normalize(torch.cat(all_embs, dim=0), p=2, dim=1)
    return qa_database, qa_embeddings


def find_similar(user_q, model, device, bpe_tok, word2idx, qa_database, qa_embeddings):
    if qa_embeddings is None:
        return None, 0.0
    sos_id = word2idx.get('<sos>', 1)
    q_ids = tokenize_bpe(user_q, bpe_tok, word2idx)
    x = torch.tensor([[sos_id] + q_ids]).to(device)
    user_emb = F.normalize(model.get_embedding(x), p=2, dim=1)
    sims = (qa_embeddings @ user_emb.T).squeeze(1)
    best_idx = sims.argmax().item()
    best_score = sims[best_idx].item()
    if best_score >= SIMILARITY_THRESHOLD:
        return qa_database[best_idx], best_score
    return None, best_score


def clean_response(text):
    text = text.strip()
    if not text:
        return ""
    # Viết hoa chữ cái đầu tiên
    text = text[0].upper() + text[1:]
    
    # Chuẩn hóa Unicode (sửa lỗi dấu tiếng Việt NFD nếu có)
    import unicodedata
    text = unicodedata.normalize('NFC', text)
    
    # Sửa lỗi khoảng trắng trong LaTeX (do Tokenizer tách)
    text = re.sub(r'\\\s+([a-zA-Z]+)', r'\\\1', text) # \ cdot -> \cdot
    text = re.sub(r'\s*_\s*', '_', text)            # k _ 5 -> k_5
    text = re.sub(r'\s*\^\s*', '^', text)            # x ^ 2 -> x^2
    text = re.sub(r'\s*\{\s*', '{', text)            # { 5 } -> {5}
    text = re.sub(r'\s*\}\s*', '}', text)
    
    # Xóa khoảng trắng thừa trước các dấu câu
    text = re.sub(r'\s+([.,!?():;])', r'\1', text)
    # Xóa khoảng trắng thừa sau dấu mở ngoặc
    text = re.sub(r'\(\s+', r'(', text)
    # Rút gọn khoảng trắng trùng lặp
    text = re.sub(r'\s{2,}', ' ', text)
    return text.strip()


def chat():
    if not os.path.exists(MODEL_PATH):
        print(f"Không tìm thấy model: {MODEL_PATH}")
        print("Hãy chạy train_gpu.py trước.")
        return

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

    with open(VOCAB_PATH, 'rb') as f:
        vocab = pickle.load(f)
    word2idx = vocab['word2idx']
    idx2word = vocab['idx2word']

    ckpt  = torch.load(MODEL_PATH, map_location=device)
    
    # Hỗ trợ load cả model.pt và checkpoint_large.pt
    d_m = ckpt.get('d_m', 768)
    n_l = ckpt.get('n_l', 12)
    n_h = ckpt.get('n_h', 12)
    dr  = ckpt.get('dr', 0.1)

    model = make_model(
        len(word2idx),
        n_layer=n_l, d_model=d_m,
        n_head=n_h,  dropout=dr
    )
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device)
    model.eval()

    bpe_tok = load_bpe_tokenizer()
    print(f"  BPE Tokenizer: {'Loaded' if bpe_tok else 'Fallback to split()'}")

    # Đã tắt để AI tự sinh 100% nhằm tăng tốc khởi động
    qa_database, qa_embeddings = [], None
    print("  [INFO] Bỏ qua build embeddings vì đã tắt Similarity Match.")

    print(f"\n--- CHATBOT TOÁN RỜI RẠC | device={device} ---")
    print("Gõ 'exit' để thoát.\n")

    while True:
        try:
            user_input = input("Bạn: ").strip()
            if user_input.lower() in ['exit', 'quit', 'thoát']:
                break
            if not user_input:
                continue

            # Bước 1: So khớp tương đồng (Đã tắt để AI tự sinh 100%)
            matched, score = None, 0.0
            if matched is not None:
                print(f"Bot: {matched['answer']}")
                print(f"  [Similarity: {score:.2%} | Matched: \"{matched['question']}\"]")
                print()
                continue

            # Bước 2: Generate
            tokens = tokenize_bpe(user_input, bpe_tok, word2idx)
            sos_id = word2idx.get('<sos>', 1)
            sep_id = word2idx.get('<sep>', 4)
            prompt = [sos_id] + tokens + [sep_id]
            x = torch.tensor([prompt]).to(device)

            with torch.no_grad():
                out_ids = model.generate(x, max_new_tokens=100,
                                         temperature=0.5, top_k=50, top_p=0.9,
                                         repetition_penalty=1.3)

            resp_ids = out_ids[0][len(prompt):]

            if bpe_tok is not None:
                valid_ids = []
                for i in resp_ids:
                    tid = i.item()
                    if tid == word2idx.get('<eos>', 2):
                        break
                    if tid not in [word2idx.get('<pad>', 0), word2idx.get('<sos>', 1),
                                   word2idx.get('<unk>', 3), word2idx.get('<sep>', 4)]:
                        valid_ids.append(tid)
                response = bpe_tok.decode(valid_ids)
            else:
                resp_tokens = []
                for i in resp_ids:
                    sym = idx2word.get(i.item(), '<unk>')
                    if sym == '<eos>':
                        break
                    if sym not in ['<pad>', '<sos>', '<unk>', '<sep>']:
                        resp_tokens.append(sym)
                response = " ".join(resp_tokens)

            # Làm sạch câu trả lời sinh ra
            response = clean_response(response)

            if not response.strip():
                response = "Mình chưa học tới phần này. Bạn hỏi câu khác nhé!"

            print(f"Bot: {response}")
            print(f"  [Generated | Similarity best: {score:.2%}]")
            print()

        except KeyboardInterrupt:
            break
        except Exception as e:
            print(f"Lỗi: {e}")


if __name__ == "__main__":
    chat()