"""
Preprocess: Dùng BPE Tokenizer để mã hóa dữ liệu.
Chạy: python src/preprocess.py
Kết quả: data/processed/data.pkl + data/processed/vocab.pkl
"""
import pandas as pd
import pickle
import os
import sys
from sklearn.model_selection import train_test_split

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

INPUT_CSV      = 'data/augmented/augmented.csv'
DATA_DIR       = 'data/processed'
VOCAB_PATH     = os.path.join(DATA_DIR, 'vocab.pkl')
DATA_PATH      = os.path.join(DATA_DIR, 'data.pkl')
TOKENIZER_PATH = os.path.join(DATA_DIR, 'bpe_tokenizer.json')

os.makedirs(DATA_DIR, exist_ok=True)


def load_tokenizer():
    """Load BPE tokenizer đã train sẵn."""
    from tokenizers import Tokenizer
    if not os.path.exists(TOKENIZER_PATH):
        print(f"Không tìm thấy {TOKENIZER_PATH}. Hãy chạy train_tokenizer.py trước.")
        return None
    return Tokenizer.from_file(TOKENIZER_PATH)


def tokenize_bpe(text, tokenizer):
    """Tokenize bằng BPE thay vì split()."""
    enc = tokenizer.encode(str(text).strip())
    return enc.ids


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

    tokenizer = load_tokenizer()
    if tokenizer is None:
        return

    vocab_size = tokenizer.get_vocab_size()
    vocab = tokenizer.get_vocab()
    # Tạo word2idx và idx2word cho tương thích với code cũ
    word2idx = vocab  # {token_str: id}
    idx2word = {v: k for k, v in vocab.items()}

    # Lưu vocab.pkl (tương thích với api.py và chat.py cũ)
    with open(VOCAB_PATH, 'wb') as f:
        pickle.dump({'word2idx': word2idx, 'idx2word': idx2word}, f)
    print(f"Vocab size: {vocab_size}")

    # Đọc CSV đã augmented
    # Tự động phát hiện separator
    with open(INPUT_CSV, 'r', encoding='utf-8-sig') as f:
        first_line = f.readline()
    sep = ';' if first_line.count(';') >= 2 else ','
    print(f"Đọc CSV với separator: '{sep}'")

    df = pd.read_csv(INPUT_CSV, sep=sep, header=None,
                     names=['topic', 'question', 'answer'],
                     encoding='utf-8-sig').dropna(subset=['question', 'answer'])

    print(f"Tổng mẫu: {len(df)}")

    # Lấy các special token IDs
    sos_id = word2idx.get('<sos>', word2idx.get('[SOS]', 1))
    eos_id = word2idx.get('<eos>', word2idx.get('[EOS]', 2))
    sep_id = word2idx.get('<sep>', word2idx.get('[SEP]', 4))

    def encode_row(row):
        q_ids = tokenize_bpe(row['question'], tokenizer)
        a_ids = tokenize_bpe(row['answer'], tokenizer)
        return [sos_id] + q_ids + [sep_id] + a_ids + [eos_id]

    # Chia train/val
    train_df, val_df = train_test_split(
        df, test_size=0.15, stratify=df['topic'], random_state=42
    )
    train_df = train_df.reset_index(drop=True)
    val_df   = val_df.reset_index(drop=True)

    train_seqs = [encode_row(row) for _, row in train_df.iterrows()]
    val_seqs   = [encode_row(row) for _, row in val_df.iterrows()]

    with open(DATA_PATH, 'wb') as f:
        pickle.dump({
            'sequences': train_seqs + val_seqs,
            'topics':    train_df['topic'].tolist() + val_df['topic'].tolist(),
            'train': {'sequences': train_seqs, 'topics': train_df['topic'].tolist()},
            'val':   {'sequences': val_seqs,   'topics': val_df['topic'].tolist()},
        }, f)

    print(f"Train: {len(train_seqs)} | Val: {len(val_seqs)}")
    print(f"Đã lưu: {DATA_PATH} và {VOCAB_PATH}")


if __name__ == "__main__":
    main()