import csv
import json
import random
import yaml
from pathlib import Path
from src.models.recurrent.tokenizer import SPMTokenizer


def _load_dataset_config():
    """Load dataset paths from config instead of hardcoding."""
    with open("config/data_config.yml", "r") as f:
        config = yaml.safe_load(f)
    return config["paths"]["seq2seq"]["datasets"]


DATASETS = _load_dataset_config()


def prepare_data_from_csv(
    csv_dir: str,
    output_dir: str,
    val_split: float = 0.1,
    vocab_size: int = 4000,
    seed: int = 42,
) -> None:
    """
    Reads each CSV in DATASETS and writes per-direction JSONL files
    to output_dir/<lang_dir>/{train,val,test}.jsonl and trains a
    SentencePiece BPE model saved as spm.model in the same directory.
    Skips if files and model already exist.
    """
    random.seed(seed)
    csv_dir = Path(csv_dir)
    output_dir = Path(output_dir)

    for lang_dir, csv_filename in DATASETS.items():
        pair_dir = output_dir / lang_dir
        spm_model_path = pair_dir / "spm.model"

        if (
            all((pair_dir / f"{s}.jsonl").exists() for s in ("train", "val", "test"))
            and spm_model_path.exists()
        ):
            print(f"{lang_dir}: already processed, skipping.")
            continue

        pair_dir.mkdir(parents=True, exist_ok=True)
        train_rows, val_rows, test_rows = [], [], []

        with open(csv_dir / csv_filename, "r", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                record = {
                    "source_text": row["source_text"].lower().strip(),
                    "target_text": row["target_text"].lower().strip(),
                }
                if row["split"] == "test":
                    test_rows.append(record)
                elif row["split"] == "val":
                    val_rows.append(record)
                elif row["split"] == "train":
                    train_rows.append(record)

        random.shuffle(train_rows)
        splits = {
            "train": train_rows,
            "val": val_rows,
            "test": test_rows,
        }
        for split_name, rows in splits.items():
            out_path = pair_dir / f"{split_name}.jsonl"
            with open(out_path, "w", encoding="utf-8") as f:
                for r in rows:
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
            print(f"{lang_dir} | {split_name}: {len(rows)} rows -> {out_path}")

        # Train a shared SentencePiece model on all source + target text from training split
        all_texts = [r["source_text"] for r in train_rows] + [
            r["target_text"] for r in train_rows
        ]
        model_prefix = str(pair_dir / "spm")
        SPMTokenizer.train(all_texts, model_prefix, vocab_size=vocab_size)
        print(f"{lang_dir}: SPM model trained -> {model_prefix}.model")