import asyncio
import json
from pathlib import Path

import pytest

from gfa_core.mas.agents.planning_v2 import PlanningAgentV2


GENOTYPE_SCOPE_ID = "genotype:" + "1" * 64
PHENOTYPE_SCOPE_ID = "phenotype:" + "2" * 64
COVARIATE_SCOPE_ID = "covariate:" + "3" * 64


def _write_manifest(final_dir: Path, stage: str, prefix: str, **scope_ids: str) -> Path:
    path = final_dir / f"{prefix}.{stage}.manifest.json"
    path.write_text(
        json.dumps(
            {
                "stage": stage,
                "prefix": prefix,
                "status": "success",
                **scope_ids,
            }
        ),
        encoding="utf-8",
    )
    return path


def _completed_preflight(workspace: Path, prefix: str = "Os") -> dict:
    final_dir = workspace / "final"
    final_dir.mkdir(parents=True)
    manifests = {
        "gwas_fmt": _write_manifest(
            final_dir,
            "gwas_fmt",
            prefix,
            genotype_scope_id=GENOTYPE_SCOPE_ID,
        ),
        "gwas_kinship": _write_manifest(
            final_dir,
            "gwas_kinship",
            prefix,
            genotype_scope_id=GENOTYPE_SCOPE_ID,
        ),
        "gwas_trait": _write_manifest(
            final_dir,
            "gwas_trait",
            prefix,
            phenotype_scope_id=PHENOTYPE_SCOPE_ID,
        ),
        "gwas_cov": _write_manifest(
            final_dir,
            "gwas_cov",
            prefix,
            genotype_scope_id=GENOTYPE_SCOPE_ID,
            covariate_scope_id=COVARIATE_SCOPE_ID,
        ),
    }
    stage_status = {
        stage: {
            "status": "success",
            "manifest_path": str(path),
            "output_files": [],
            "missing_outputs": [],
        }
        for stage, path in manifests.items()
    }
    stage_status["gwas_mlm"] = {
        "status": "not_found",
        "manifest_path": None,
        "output_files": [],
        "missing_outputs": [],
    }
    return {"stage_status": stage_status}


def _plan_mlm(workspace: Path, preflight_result: dict) -> dict:
    agent = PlanningAgentV2.__new__(PlanningAgentV2)
    intent = {
        "stages": ["gwas_mlm"],
        "prefix": "Os",
        "params": {},
        "task_kind": "pipeline",
    }
    return asyncio.run(
        agent._generate_plan_with_llm(
            "使用 TASSEL 执行 GWAS MLM 分析",
            intent,
            completed_stages=[],
            job_id="job_test",
            workspace=str(workspace),
            preflight_result=preflight_result,
        )
    )


def test_completed_gwas_dependencies_handoff_exact_manifest_scopes_to_mlm(tmp_path: Path) -> None:
    workspace = tmp_path / "workspace" / "job_test"
    plan = _plan_mlm(workspace, _completed_preflight(workspace))

    assert [stage["id"] for stage in plan["stages"]] == ["gwas_mlm"]
    params = plan["stages"][0]["params"]
    assert params["genotypeScopeId"] == GENOTYPE_SCOPE_ID
    assert params["phenotypeScopeId"] == PHENOTYPE_SCOPE_ID
    assert params["covariateScopeId"] == COVARIATE_SCOPE_ID


def _rewrite_manifest(path: Path, **updates: str) -> None:
    manifest = json.loads(path.read_text(encoding="utf-8"))
    manifest.update(updates)
    path.write_text(json.dumps(manifest), encoding="utf-8")


@pytest.mark.parametrize(
    "invalid_case",
    ["missing_manifest", "wrong_prefix", "illegal_scope", "genotype_conflict"],
)
def test_completed_gwas_dependency_identity_is_all_or_nothing(
    tmp_path: Path,
    invalid_case: str,
) -> None:
    workspace = tmp_path / "workspace" / "job_test"
    preflight = _completed_preflight(workspace)
    status = preflight["stage_status"]
    if invalid_case == "missing_manifest":
        status["gwas_trait"]["manifest_path"] = None
    elif invalid_case == "wrong_prefix":
        _rewrite_manifest(Path(status["gwas_trait"]["manifest_path"]), prefix="At")
    elif invalid_case == "illegal_scope":
        _rewrite_manifest(Path(status["gwas_trait"]["manifest_path"]), phenotype_scope_id="bad")
    else:
        _rewrite_manifest(
            Path(status["gwas_cov"]["manifest_path"]),
            genotype_scope_id="genotype:" + "4" * 64,
        )

    params = _plan_mlm(workspace, preflight)["stages"][0]["params"]

    assert "genotypeScopeId" not in params
    assert "phenotypeScopeId" not in params
    assert "covariateScopeId" not in params