import pandas as pd
import pytest

from giant.data.loader import build_process_map_from_files, find_parquet_files


def _touch(path):
    path.parent.mkdir(parents=True, exist_ok=True)
    path.touch()
    return path


def test_find_parquet_files_single_file(tmp_path):
    f = _touch(tmp_path / "shard-000.parquet")
    assert find_parquet_files(f) == [f]


def test_find_parquet_files_directory_glob(tmp_path):
    a = _touch(tmp_path / "shard-000.parquet")
    b = _touch(tmp_path / "shard-001.parquet")
    _touch(tmp_path / "not_a_parquet.root")
    assert find_parquet_files(tmp_path) == sorted([a, b])


def test_find_parquet_files_empty_directory_raises(tmp_path):
    with pytest.raises(FileNotFoundError):
        find_parquet_files(tmp_path)


def test_manifest_resolves_relative_to_its_own_directory(tmp_path):
    target = _touch(tmp_path / "processed" / "pbwo4" / "shard-000.parquet")
    manifest_dir = tmp_path / "pools" / "pbwo4"
    manifest_dir.mkdir(parents=True)
    manifest = manifest_dir / "full.manifest"
    manifest.write_text("../../processed/pbwo4/shard-000.parquet\n")

    assert find_parquet_files(manifest) == [target.resolve()]


def test_manifest_skips_blank_lines_and_comments(tmp_path):
    target = _touch(tmp_path / "shard-000.parquet")
    manifest = tmp_path / "full.manifest"
    manifest.write_text("\n# a comment\nshard-000.parquet\n\n")

    assert find_parquet_files(manifest) == [target.resolve()]


def test_manifest_missing_file_raises(tmp_path):
    manifest = tmp_path / "full.manifest"
    manifest.write_text("does_not_exist.parquet\n")

    with pytest.raises(FileNotFoundError):
        find_parquet_files(manifest)


def test_manifest_with_no_entries_raises(tmp_path):
    manifest = tmp_path / "full.manifest"
    manifest.write_text("# only comments\n")

    with pytest.raises(FileNotFoundError):
        find_parquet_files(manifest)


def test_build_process_map_from_files_keeps_most_frequent(tmp_path):
    """process counts: eIoni=5, phot=3, compt=2, Rayl=1 — with n_experts=3, only
    the top 2 (eIoni, phot) get their own index; compt/Rayl share the "other"
    (last) index."""
    process = ["eIoni"] * 5 + ["phot"] * 3 + ["compt"] * 2 + ["Rayl"] * 1
    path = tmp_path / "shard-000.parquet"
    pd.DataFrame({"process": process}).to_parquet(path)

    proc_map = build_process_map_from_files([path], n_experts=3)

    assert proc_map["eIoni"] == 0
    assert proc_map["phot"] == 1
    assert proc_map["compt"] == 2
    assert proc_map["Rayl"] == 2
    assert set(proc_map.values()) <= {0, 1, 2}


def test_build_process_map_from_files_spans_multiple_files(tmp_path):
    path_a = tmp_path / "a.parquet"
    path_b = tmp_path / "b.parquet"
    pd.DataFrame({"process": ["eIoni"] * 3 + ["phot"] * 1}).to_parquet(path_a)
    pd.DataFrame({"process": ["phot"] * 4 + ["compt"] * 1}).to_parquet(path_b)

    # phot: 1+4=5 total > eIoni: 3 > compt: 1
    proc_map = build_process_map_from_files([path_a, path_b], n_experts=3)

    assert proc_map["phot"] == 0
    assert proc_map["eIoni"] == 1
    assert proc_map["compt"] == 2