import json
import socket
import subprocess
import sys
from pathlib import Path

import pytest
import torch

from mouh.training import PackedByteDataset, SFTDataset, TrainArgs, lr_at, train
from mouh.verification.correction_store import CorrectionStore
from mouh.jev.types import JevDiagnostic, JevErrorType, JevSeverity

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


def corpus(tmp_path):
    f = tmp_path / "c.txt"
    f.write_text(("def add(a, b):\n    return a + b\n\n" * 40))
    return str(f)


def test_packed_dataset_blocks_are_disjoint_between_ranks(tmp_path):
    texts = ["x" * 500]  # un seul document : le partage par rang doit quand même fonctionner
    r0 = list(PackedByteDataset(texts, 64, rank=0, world_size=2, loop=False))
    r1 = list(PackedByteDataset(texts, 64, rank=1, world_size=2, loop=False))
    assert len(r0) > 0 and len(r1) > 0 and abs(len(r0) - len(r1)) <= 1
    assert all(b.shape == (64,) for b in r0 + r1)


def test_lr_schedule():
    a = TrainArgs(lr=1.0, warmup_steps=10, max_steps=110, min_lr_ratio=0.1)
    assert lr_at(0, a) == pytest.approx(0.1) and lr_at(9, a) == pytest.approx(1.0)
    assert lr_at(110, a) == pytest.approx(0.1, abs=1e-6)
    assert lr_at(60, a) < lr_at(20, a)


def test_single_process_training_reduces_loss_and_saves_loadable_checkpoint(tmp_path):
    a = TrainArgs(preset="mouh-tiny", data=corpus(tmp_path), seq_len=96, micro_batch_size=4, grad_accum=1,
                  max_steps=25, warmup_steps=3, lr=3e-3, log_every=100, save_every=0,
                  output_dir=str(tmp_path / "out"), device="cpu")
    res = train(a)
    assert res["last_loss"] < res["first_loss"] * 0.7
    from mouh.model import MouhForCausalLM
    m = MouhForCausalLM.from_pretrained(tmp_path / "out" / "final")
    assert m(torch.randint(0, 256, (1, 8))).logits.shape[-1] == 512
    assert json.loads((tmp_path / "out" / "final" / "trainer_state.json").read_text())["step"] == 25


def test_sft_from_correction_store(tmp_path):
    store = CorrectionStore(tmp_path / "c.jsonl")
    d = JevDiagnostic(severity=JevSeverity.ERROR, error_type=JevErrorType.TEST_FAILURE, message="assert 1 == 2")
    store.add(task="add", failing_files={"a.py": "x"}, diagnostics=(d,), fixed_files={"a.py": "y"}, attempts=2)
    store.export_sft(tmp_path / "sft.jsonl")
    ds = SFTDataset(str(tmp_path / "sft.jsonl"), seq_len=512)
    ids, labels = next(iter(ds))
    assert ids.shape == labels.shape == (512,)
    assert (labels == -100).sum() > 0 and (labels != -100).sum() > 0   # perte sur la complétion seulement


def _free_port():
    with socket.socket() as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


@pytest.mark.slow
def test_ddp_two_processes_cpu(tmp_path):
    cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", f"--master_port={_free_port()}",
           str(ROOT / "scripts" / "train.py"), "--preset", "mouh-tiny", "--data", corpus(tmp_path), "--seq_len", "96",
           "--micro_batch_size", "2", "--grad_accum", "2", "--max_steps", "6", "--warmup_steps", "2", "--lr", "3e-3",
           "--log_every", "3", "--save_every", "0", "--parallel", "ddp", "--device", "cpu",
           "--output_dir", str(tmp_path / "ddp")]
    p = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    assert p.returncode == 0, p.stderr[-2000:]
    assert (tmp_path / "ddp" / "final" / "model.safetensors").exists()


def test_text_source_is_reiterable_so_epochs_loop(tmp_path):
    from mouh.training import TextSource
    ds = PackedByteDataset(TextSource(corpus(tmp_path)), 96, loop=True)
    it = iter(ds)
    blocks = [next(it) for _ in range(60)]   # bien plus qu'une époque (~15 blocs)
    assert len(blocks) == 60