from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pytest
import yaml

REPO = Path(__file__).resolve().parents[1]

# Texto en español sintético pero con la estructura correcta (acentos, ñ,
# signos de apertura): suficiente para los tests, y sin depender de la red.
_MUESTRA = (
    "ENLACE responde de forma concreta. No adorna. La luz del living está "
    "encendida y la temperatura del cuarto es de veintiún grados. ¿Querés que "
    "apague la del pasillo? El calendario tiene una entrada mañana temprano. "
    "Andrew aprendió despacio, un día a la vez, y así fue construyendo quién "
    "era. La memoria no está en los pesos: está en la base de datos. "
)


@pytest.fixture(scope="session")
def texto_es(tmp_path_factory: pytest.TempPathFactory) -> Path:
    path = tmp_path_factory.mktemp("datos") / "texto.txt"
    path.write_text(_MUESTRA * 200, encoding="utf-8")
    return path


@pytest.fixture(scope="session")
def shards_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
    """Dos shards uint16 con tokens predecibles, para verificar el recorrido."""
    directory = tmp_path_factory.mktemp("shards")
    total = 0
    entries = []
    for i, count in enumerate([4096, 4096]):
        tokens = np.arange(total, total + count, dtype=np.uint16)
        name = f"shard-{i:04d}.bin"
        tokens.tofile(directory / name)
        entries.append({"file": name, "tokens": count})
        total += count
    (directory / "index.json").write_text(
        json.dumps({"vocab_size": 8192, "shards": entries, "total_tokens": total})
    )
    return directory


def write_run_config(
    tmp_path: Path,
    texto: Path,
    *,
    run_name: str,
    max_steps: int,
    checkpoint_every: int,
) -> Path:
    """Config de corrida mínima en CPU, apuntando a un texto de prueba."""
    cfg = {
        "hardware": yaml.safe_load((REPO / "configs/hardware/cpu.yaml").read_text()),
        "model": yaml.safe_load((REPO / "configs/model/char-smoke.yaml").read_text()),
        "train": {
            "run_name": run_name,
            "out_dir": str(tmp_path / "runs"),
            "seed": 1337,
            "max_steps": max_steps,
            "optimizer": {
                "lr": 3.0e-3,
                "beta1": 0.9,
                "beta2": 0.95,
                "eps": 1.0e-8,
                "weight_decay": 0.1,
                "grad_clip": 1.0,
            },
            "schedule": {
                "kind": "wsd",
                "warmup_steps": 2,
                "decay_steps": 2,
                "min_lr_ratio": 0.0,
            },
            "log_every": 1000,
            "eval_every": 1000,
            "eval_batches": 2,
            "checkpoint_every": checkpoint_every,
            "sample_every": 0,
        },
        "data": {
            "source": "chars",
            "text_path": str(texto),
            "val_fraction": 0.05,
        },
    }
    # Modelo aún más chico: los tests tienen que correr en segundos.
    cfg["model"].update({"n_layer": 2, "d_model": 64, "n_head": 4, "n_kv_head": 2, "seq_len": 64})
    cfg["hardware"].update({"micro_batch_size": 4, "grad_accum_steps": 2})

    path = tmp_path / f"{run_name}.yaml"
    path.write_text(yaml.safe_dump(cfg))
    return path