from __future__ import annotations

from dataclasses import fields
from pathlib import Path

from gfa_core.mas.stage_capabilities import CAPABILITIES
from gfa_core.mas.stage_catalog import (
    EXECUTABLE_STAGES,
    PIPELINE_STAGES,
    STAGE_DEPENDENCIES,
    STAGE_INFO,
    STAGE_RESOURCE_PROFILES,
    complete_dependencies,
)
from gfa_core.mas.tools.stage_identity import StageIdentity, params_from_identity
from gfa_core.mas.tools.stage_tools import StageToolRegistry


PROJECT_ROOT = Path(__file__).resolve().parents[2]

GWAS_STAGES = {
    "gwas_kinship": "run_gwas_kinship.py",
    "gwas_trait": "run_gwas_trait.py",
    "gwas_cov": "run_gwas_covariate.py",
    "gwas_mlm": "run_gwas_tassel_mlm.py",
    "gwas_plot": "run_gwas_plot.py",
    "gwas_emmax": "run_gwas_emmax.py",
    "gwas_emplot": "run_gwas_emmax_plot.py",
    "gwas_glm": "run_gwas_tassel_glm.py",
}

EXPECTED_DEPENDENCIES = {
    "gwas_kinship": ["gwas_fmt"],
    "gwas_trait": [],
    "gwas_cov": [],
    "gwas_mlm": ["gwas_fmt", "gwas_kinship", "gwas_trait", "gwas_cov"],
    "gwas_plot": ["gwas_mlm"],
    "gwas_emmax": ["gwas_fmt", "gwas_trait"],
    "gwas_emplot": ["gwas_emmax"],
    "gwas_glm": ["gwas_fmt", "gwas_trait", "gwas_cov"],
}


def test_all_gwas_stages_have_one_stable_registration_and_runner() -> None:
    definitions = {item["id"]: item for item in StageToolRegistry.STAGE_DEFINITIONS}
    for stage, runner in GWAS_STAGES.items():
        assert stage in PIPELINE_STAGES
        assert stage in EXECUTABLE_STAGES
        assert stage in STAGE_INFO
        assert stage in CAPABILITIES
        assert definitions[stage]["script"] == runner
        assert CAPABILITIES[stage].cli == f"apps/cli/{runner}"


def test_gwas_dependencies_are_exact_and_acyclic() -> None:
    for stage, dependencies in EXPECTED_DEPENDENCIES.items():
        assert STAGE_DEPENDENCIES[stage] == dependencies
    full_chain = complete_dependencies(["gwas_plot", "gwas_emplot", "gwas_glm"])
    assert len(full_chain) == len(set(full_chain))
    for stage in full_chain:
        for dependency in STAGE_DEPENDENCIES.get(stage, []):
            assert full_chain.index(dependency) < full_chain.index(stage)


def test_model_stages_share_heavy_capacity_group() -> None:
    for stage in ("gwas_mlm", "gwas_emmax", "gwas_glm"):
        profile = STAGE_RESOURCE_PROFILES[stage]
        assert profile["queue"] == "heavy"
        assert profile["resource_class"] == "heavy"
    for stage in ("gwas_trait", "gwas_cov", "gwas_plot", "gwas_emplot"):
        assert STAGE_RESOURCE_PROFILES[stage]["resource_class"] == "light"


def test_gwas_identity_fields_round_trip_to_api_params() -> None:
    names = {field.name for field in fields(StageIdentity)}
    required = {
        "genotype_scope_id",
        "phenotype_scope_id",
        "covariate_scope_id",
        "model_scope_id",
        "source_model_scope_id",
        "analysis_scope_id",
        "run_id",
    }
    assert required <= names
    identity = StageIdentity(
        stage="gwas_mlm",
        prefix="cohort",
        genotype_scope_id="genotype:g",
        phenotype_scope_id="phenotype:p",
        covariate_scope_id="covariate:c",
        model_scope_id="gwas-model:m",
        source_model_scope_id="gwas-model:s",
        analysis_scope_id="gwas-analysis:a",
        run_id="gwas_mlm_run",
    )
    params = params_from_identity(identity)
    assert params["genotypeScopeId"] == "genotype:g"
    assert params["phenotypeScopeId"] == "phenotype:p"
    assert params["covariateScopeId"] == "covariate:c"
    assert params["modelScopeId"] == "gwas-model:m"
    assert params["sourceModelScopeId"] == "gwas-model:s"
    assert params["analysisScopeId"] == "gwas-analysis:a"
    assert params["runId"] == "gwas_mlm_run"


def test_each_gwas_stage_has_contract_and_skill_boundary() -> None:
    skill_names = {
        "gwas_kinship": "gwas-kinship",
        "gwas_trait": "gwas-trait",
        "gwas_cov": "gwas-covariate",
        "gwas_mlm": "gwas-tassel-mlm",
        "gwas_plot": "gwas-plot",
        "gwas_emmax": "gwas-emmax",
        "gwas_emplot": "gwas-emmax-plot",
        "gwas_glm": "gwas-tassel-glm",
    }
    for stage, skill in skill_names.items():
        assert (PROJECT_ROOT / "gfa_core" / "schemas" / "stages" / f"{stage}_contract_v1.md").is_file()
        assert (PROJECT_ROOT / "skills" / skill / "SKILL.md").is_file()
        assert (PROJECT_ROOT / "skills" / skill / "config.yaml").is_file()
        assert (PROJECT_ROOT / "skills" / skill / "run.py").is_file()