from __future__ import annotations

import json
from pathlib import Path

import pytest

from deepgenopix.config import ExperimentConfig
from deepgenopix.notebook_support import (
    NotebookPaths,
    TrainNotebookContext,
    ensure_train_artifacts,
    prepare_train_context,
)


def _make_context(tmp_path: Path, *, run_etl: bool = True):
    repo_root = tmp_path / "repo"
    run_root = repo_root / "data" / "output" / "baseline_v1" / "baseline_v1__px12"
    processed_root = repo_root / "data" / "processed" / "te_visuals" / "baseline_v1__px12"
    raw_root = repo_root / "data" / "raw"
    raw_fasta = raw_root / "input.fa"

    for directory in [run_root, processed_root, raw_root]:
        directory.mkdir(parents=True, exist_ok=True)
    raw_fasta.write_text(">seq0\nACGTACGTACGT\n", encoding="utf-8")

    paths = NotebookPaths(
        repo_root=repo_root,
        run_root=run_root,
        processed_root=processed_root,
        checkpoint_dir=run_root / "checkpoints",
        report_dir=run_root / "reports",
        embedding_dir=run_root / "embeddings",
        config_path=run_root / "config.json",
        metrics_path=run_root / "metrics.csv",
        manifest_path=run_root / "run_manifest.json",
        notes_path=run_root / "run_notes.md",
        report_path=run_root / "run_report.md",
        raw_input_dir=raw_root,
        raw_fasta=raw_fasta,
        raw_parquet=None,
        raw_split_parquet_root=None,
        raw_split_summary=None,
    )

    return TrainNotebookContext(
        config=ExperimentConfig(),
        paths=paths,
        device="cpu",
        in_colab=False,
        run_mode="train",
        resume_from_existing=False,
        run_etl=run_etl,
        copy_data_manifest=False,
        export_embeddings=False,
    )


def test_ensure_train_artifacts_rebuilds_partial_dataset(monkeypatch, tmp_path):
    context = _make_context(tmp_path)
    processed_root = context.paths.processed_root
    run_root = context.paths.run_root

    (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8")
    (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8")

    calls: list[str] = []

    def fake_run_incremental_etl(fasta_path, output_dir, **kwargs):
        calls.append(str(fasta_path))
        assert not (processed_root / "registry.csv").exists()
        assert not (processed_root / "classes.json").exists()
        assert not (processed_root / "tensors.lmdb").exists()
        (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8")
        (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8")
        (processed_root / "tensors.lmdb").write_bytes(b"lmdb")

    monkeypatch.setattr("deepgenopix.notebook_support.run_incremental_etl", fake_run_incremental_etl)

    artifacts = ensure_train_artifacts(context)

    assert calls == [str(context.paths.raw_fasta)]
    assert artifacts.registry_path == processed_root / "registry.csv"
    assert artifacts.classes_path == processed_root / "classes.json"
    assert (processed_root / "tensors.lmdb").exists()
    assert (run_root / "run_manifest.json").exists()


def test_ensure_train_artifacts_raises_when_lmdb_missing_and_reuse_only(monkeypatch, tmp_path):
    context = _make_context(tmp_path, run_etl=False)
    processed_root = context.paths.processed_root

    (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8")
    (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8")

    called = False

    def fake_run_incremental_etl(*args, **kwargs):
        nonlocal called
        called = True

    monkeypatch.setattr("deepgenopix.notebook_support.run_incremental_etl", fake_run_incremental_etl)

    try:
        ensure_train_artifacts(context)
    except FileNotFoundError as exc:
        assert "tensors.lmdb" in str(exc)
    else:
        raise AssertionError("expected FileNotFoundError")

    assert called is False


def test_ensure_train_artifacts_uses_pre_split_parquets(monkeypatch, tmp_path):
    context = _make_context(tmp_path)
    processed_root = context.paths.processed_root
    split_root = context.paths.raw_input_dir
    context.paths.raw_fasta = None
    context.paths.raw_split_parquet_root = split_root

    for split in ("train", "val", "test"):
        split_dir = split_root / split
        split_dir.mkdir(parents=True, exist_ok=True)
        (split_dir / "te_seqdata.parquet").write_bytes(b"parquet")

    calls: list[str] = []

    def fake_run_split_parquet_etl(split_input_root, output_dir, **kwargs):
        calls.append(str(split_input_root))
        (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8")
        (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8")
        (processed_root / "tensors.lmdb").write_bytes(b"lmdb")

    monkeypatch.setattr("deepgenopix.notebook_support.run_split_parquet_etl", fake_run_split_parquet_etl)

    artifacts = ensure_train_artifacts(context)

    assert calls == [str(split_root)]
    assert artifacts.registry_path == processed_root / "registry.csv"


def test_prepare_train_context_refuses_to_overwrite_completed_run(tmp_path):
    repo_root = tmp_path / "repo"
    raw_root = repo_root / "data" / "raw"
    run_root = repo_root / "data" / "output" / "baseline_v1" / ExperimentConfig().run_id
    raw_root.mkdir(parents=True, exist_ok=True)
    run_root.mkdir(parents=True, exist_ok=True)
    (raw_root / "input.fa").write_text(">seq0\nACGTACGTACGT\n", encoding="utf-8")
    (run_root / "metrics.csv").write_text("epoch,train_loss,val_accuracy\n1,1.0,0.5\n", encoding="utf-8")
    (run_root / "run_manifest.json").write_text(json.dumps({"status": "complete"}), encoding="utf-8")

    with pytest.raises(FileExistsError, match="Run outputs already exist"):
        prepare_train_context(
            ExperimentConfig(),
            run_mode="train",
            resume_from_existing=False,
            run_etl=True,
            copy_data_manifest=False,
            export_embeddings=False,
            repo_root=repo_root,
            raw_input_dir=raw_root,
        )


def test_prepare_train_context_uses_pixel_signature_processed_root_and_split_summary(tmp_path):
    repo_root = tmp_path / "repo"
    raw_root = repo_root / "data" / "raw"

    for split in ("train", "val", "test"):
        split_dir = raw_root / split
        split_dir.mkdir(parents=True, exist_ok=True)
        (split_dir / "te_seqdata.parquet").write_bytes(b"parquet")
    (raw_root / "split_summary.json").write_text("{}", encoding="utf-8")

    context = prepare_train_context(
        ExperimentConfig(pixel_stride_bp=9, val_frac=0.2, test_frac=0.1, split_seed=7, min_family_size=2),
        run_mode="train",
        resume_from_existing=False,
        run_etl=True,
        copy_data_manifest=False,
        export_embeddings=False,
        repo_root=repo_root,
        raw_input_dir=raw_root,
    )

    assert context.paths.processed_root == (
        repo_root / "data" / "processed" / "te_visuals" / "px12_ps9_fl500-500_vf0p2_tf0p1_seed7_mf2"
    )
    assert context.paths.raw_split_parquet_root == raw_root
    assert context.paths.raw_split_summary == raw_root / "split_summary.json"
    manifest = json.loads(context.paths.manifest_path.read_text(encoding="utf-8"))
    assert manifest["raw_split_summary"] == str(raw_root / "split_summary.json")


def test_training_loop_defers_amp_choice_to_trainer_env():
    source = Path("src/deepgenopix/notebook_support.py").read_text()
    assert "use_amp=None" in source
    assert 'use_amp=str(context.device).startswith("cuda")' not in source