from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
from threading import Barrier, Thread
import zipfile

import pysam
import pytest


SAMPLES = "wgs-samples:" + "a" * 64
REFERENCE = "wgs-reference:" + "b" * 64
FILTER_ANALYSIS = "wgs-analysis:" + "c" * 64
SITE_KEYS = [
    ("chr1", 10, "A", "C"),
    ("chr1", 20, "G", "T"),
]


def _record(workspace: Path, path: Path, role: str) -> dict:
    return {
        "artifact": role,
        "role": role,
        "path": str(path.relative_to(workspace)),
        "full_path": str(path.resolve()),
        "size": path.stat().st_size,
        "size_bytes": path.stat().st_size,
        "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
    }


def _write_bgzip_vcf(
    path: Path,
    *,
    samples: tuple[str, ...] = ("102", "108", "112"),
    site_keys: list[tuple[str, int, str, str]] | None = None,
) -> None:
    plain = path.with_suffix("").with_suffix("")
    with plain.open("w", encoding="utf-8") as handle:
        handle.write("##fileformat=VCFv4.2\n")
        handle.write("##contig=<ID=chr1,length=1000>\n")
        handle.write(
            '##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\n'
        )
        handle.write(
            "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t"
            + "\t".join(samples)
            + "\n"
        )
        for chrom, pos, ref, alt in site_keys or SITE_KEYS:
            handle.write(
                f"{chrom}\t{pos}\t.\t{ref}\t{alt}\t50\tPASS\t.\tGT\t"
                + "\t".join("./." if index == 1 else "0/1" for index, _ in enumerate(samples))
                + "\n"
            )
    pysam.tabix_compress(str(plain), str(path), force=True)
    pysam.tabix_index(str(path), preset="vcf", force=True)
    plain.unlink()


def _filter_authority(
    workspace: Path,
    *,
    method: str = "gatk",
    samples: tuple[str, ...] = ("102", "108", "112"),
) -> Path:
    bundle = (
        workspace
        / "final"
        / ".wgs_variant_filter_runs"
        / "Cohort"
        / method
        / ("8" * 32)
    )
    bundle.mkdir(parents=True)
    vcf = bundle / "Cohort.wgs_variant_filter.cohort.pass.snp.vcf.gz"
    _write_bgzip_vcf(vcf, samples=samples)
    index = Path(f"{vcf}.tbi")
    report = bundle / "Cohort.wgs_variant_filter.cohort.report.md"
    report.write_text(f"method: {method}\n", encoding="utf-8")
    authority = (
        workspace
        / "final"
        / f"Cohort.wgs_variant_filter.{method}.manifest.json"
    )
    authority.write_text(
        json.dumps(
            {
                "manifest_version": "1.0",
                "stage": "wgs_variant_filter",
                "status": "success",
                "prefix": "Cohort",
                "job_id": "job-1",
                "run_id": "8" * 32,
                "workspace": str(workspace.resolve()),
                "wgs_sample_set_id": SAMPLES,
                "wgs_reference_id": REFERENCE,
                "wgs_analysis_id": FILTER_ANALYSIS,
                "method": method,
                "parameters": {"method": method},
                "metrics": {
                    "input_sample_ids": list(samples),
                    "output_sample_ids": list(samples),
                },
                "artifacts": {
                    "pass_snp_vcf": _record(
                        workspace, vcf, "pass_snp_vcf"
                    ),
                    "pass_snp_vcf_index": _record(
                        workspace, index, "pass_snp_vcf_index"
                    ),
                    "variant_filter_report": _record(
                        workspace, report, "variant_filter_report"
                    ),
                },
            }
        ),
        encoding="utf-8",
    )
    return authority


def _java(tmp_path: Path, version: str = "23-test") -> Path:
    java = tmp_path / "java"
    java.write_text(
        """#!/home/xia-lab/miniconda3/envs/test4py311/bin/python
import gzip, json, os, pathlib, sys, tempfile
import pysam
args = sys.argv[1:]
if args == ['-version']:
    print('openjdk version "JAVA_VERSION"', file=sys.stderr)
    raise SystemExit(0)
calls = os.environ.get('WGS_IMPUTE_CALLS')
if calls:
    with open(calls, 'a', encoding='utf-8') as handle:
        handle.write(json.dumps(args) + '\\n')
if os.environ.get('FAIL_POINT') == 'beagle':
    raise SystemExit(23)
def value(key):
    return next(item.split('=', 1)[1] for item in args if item.startswith(key + '='))
source = pathlib.Path(value('gt'))
out = pathlib.Path(value('out'))
with gzip.open(source, 'rt', encoding='utf-8') as handle:
    lines = list(handle)
mode = os.environ.get('BEAGLE_OUTPUT', '')
if mode == 'empty':
    lines = []
elif mode == 'invalid':
    lines = ['not-a-vcf\\n']
elif mode == 'sample-mismatch':
    lines = [line.replace('\\t112\\n', '\\t999\\n') if line.startswith('#CHROM') else line for line in lines]
elif mode == 'site-drop':
    records = [line for line in lines if not line.startswith('#')]
    lines = [line for line in lines if line.startswith('#')] + records[:1]
elif mode == 'site-reorder':
    records = [line for line in lines if not line.startswith('#')]
    lines = [line for line in lines if line.startswith('#')] + list(reversed(records))
plain = pathlib.Path(str(out) + '.vcf')
plain.parent.mkdir(parents=True, exist_ok=True)
with plain.open('w', encoding='utf-8') as handle:
    for line in lines:
        if not line.startswith('#') and mode not in ('empty', 'invalid'):
            fields = line.rstrip().split('\\t')
            fields[9:] = [gt.replace('/', '|').replace('.', '0') for gt in fields[9:]]
            line = '\\t'.join(fields) + '\\n'
        handle.write(line)
pysam.tabix_compress(str(plain), str(out) + '.vcf.gz', force=True)
plain.unlink()
pathlib.Path(str(out) + '.log').write_text('beagle test log\\n', encoding='utf-8')
""".replace("JAVA_VERSION", version),
        encoding="utf-8",
    )
    java.chmod(0o755)
    return java


def _class_jar(tmp_path: Path, class_major: int) -> Path:
    jar = tmp_path / f"beagle-java-{class_major - 44}.jar"
    class_header = (
        b"\xca\xfe\xba\xbe"
        + (0).to_bytes(2, "big")
        + class_major.to_bytes(2, "big")
    )
    with zipfile.ZipFile(jar, "w") as archive:
        archive.writestr("main/Main.class", class_header)
    return jar


def _argv(workspace: Path, authority: Path, java: Path) -> list[str]:
    return [
        "--workspace",
        str(workspace),
        "--job-id",
        "job-1",
        "--prefix",
        "Cohort",
        "--filter-manifest",
        str(authority),
        "--java",
        str(java),
        "--beagle-jar",
        str(Path("/tmp/beagle.29Oct24.c8e.jar")),
    ]


def _authority_path(workspace: Path) -> Path:
    return workspace / "final" / "Cohort.wgs_impute.manifest.json"


def _failed(workspace: Path) -> dict:
    paths = sorted((workspace / "final").glob("*wgs_impute*failed*.json"))
    if paths:
        path = paths[-1]
    else:
        path = _authority_path(workspace)
    payload = json.loads(path.read_text(encoding="utf-8"))
    assert payload["status"] == "failed"
    return payload


def test_cli_success_locks_parameters_lineage_roles_samples_sites_and_index(
    tmp_path: Path, monkeypatch
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / "workspace"
    authority = _filter_authority(workspace, method="bcftools")
    java = _java(tmp_path)
    calls = tmp_path / "calls.jsonl"
    monkeypatch.setenv("WGS_IMPUTE_CALLS", str(calls))
    assert main(_argv(workspace, authority, java)) == 0
    manifest = json.loads(_authority_path(workspace).read_text())
    assert set(manifest["artifacts"]) == {
        "imputed_snp_vcf",
        "imputed_snp_vcf_index",
        "impute_report",
        "impute_log",
    }
    assert manifest["upstream_wgs_analysis_ids"] == {
        "wgs_variant_filter": FILTER_ANALYSIS
    }
    assert manifest["upstream_method"] == "bcftools"
    assert manifest["wgs_sample_set_id"] == SAMPLES
    assert manifest["wgs_reference_id"] == REFERENCE
    assert manifest["metrics"]["input_sample_ids"] == ["102", "108", "112"]
    assert manifest["metrics"]["output_sample_ids"] == ["102", "108", "112"]
    assert manifest["metrics"]["input_site_keys"] == [
        ["chr1", 10, "A", "C"],
        ["chr1", 20, "G", "T"],
    ]
    assert manifest["metrics"]["output_site_keys"] == manifest["metrics"]["input_site_keys"]
    assert manifest["parameters"] == {
        "method": "beagle",
        "java": str(java),
        "beagle_jar": "/tmp/beagle.29Oct24.c8e.jar",
        "memory": "8g",
        "window": 40.0,
        "overlap": 2.0,
        "ne": 100000,
        "iterations": 12,
        "phase_states": 280,
        "imp_states": 1600,
        "cluster": 0.005,
    }
    command = json.loads(calls.read_text().splitlines()[0])
    assert command == [
        "-Xmx8g",
        "-jar",
        "/tmp/beagle.29Oct24.c8e.jar",
        f"gt={json.loads(authority.read_text())['artifacts']['pass_snp_vcf']['full_path']}",
        "out=" + str(workspace / ".attempts" / manifest["run_id"] / "beagle"),
        "window=40.0",
        "overlap=2.0",
        "ne=100000",
        "iterations=12",
        "phase-states=280",
        "imp-states=1600",
        "cluster=0.005",
    ]
    output = Path(manifest["artifacts"]["imputed_snp_vcf"]["full_path"])
    with pysam.VariantFile(str(output)) as handle:
        assert [record.pos for record in handle.fetch("chr1", 0, 1000)] == [10, 20]


def test_cli_rejects_java_older_than_beagle_class_version(
    tmp_path: Path
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / "workspace"
    java = _java(tmp_path, version="11.0.13")
    argv = _argv(workspace, _filter_authority(workspace), java)
    argv[-1] = str(_class_jar(tmp_path, class_major=65))

    assert main(argv) == 1
    error = _failed(workspace)["errors"][0]
    assert "requires Java 21" in error
    assert "configured Java is 11" in error


@pytest.mark.parametrize("method", ["gatk", "bcftools"])
def test_cli_accepts_either_legal_stage10_method(
    tmp_path: Path, method: str
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / method
    assert main(
        _argv(workspace, _filter_authority(workspace, method=method), _java(tmp_path))
    ) == 0


@pytest.mark.parametrize(
    ("failure", "message"),
    [
        ("beagle", "Beagle failed with exit code 23"),
        ("sample-mismatch", "sample header mismatch"),
        ("site-drop", "site keys are not conserved"),
        ("site-reorder", "building of index"),
        ("empty", "invalid"),
        ("invalid", "invalid"),
    ],
)
def test_cli_rejects_tool_and_postflight_failures_with_empty_artifacts(
    tmp_path: Path, monkeypatch, failure: str, message: str
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / "workspace"
    upstream = _filter_authority(workspace)
    if failure == "beagle":
        monkeypatch.setenv("FAIL_POINT", failure)
    else:
        monkeypatch.setenv("BEAGLE_OUTPUT", failure)
    assert main(_argv(workspace, upstream, _java(tmp_path))) == 1
    payload = _failed(workspace)
    assert payload["artifacts"] == {}
    assert message.lower() in payload["errors"][0].lower()


def test_cli_index_failure_is_not_overwritten_by_later_validation(
    tmp_path: Path, monkeypatch
) -> None:
    from apps.cli import run_wgs_impute

    workspace = tmp_path / "workspace"
    upstream = _filter_authority(workspace)

    def fail_index(*_args, **_kwargs):
        raise OSError("controlled index failure")

    monkeypatch.setattr(run_wgs_impute, "_tabix_index", fail_index)
    assert run_wgs_impute.main(_argv(workspace, upstream, _java(tmp_path))) == 1
    payload = _failed(workspace)
    assert payload["artifacts"] == {}
    assert "controlled index failure" in payload["errors"][0]


def test_failed_rerun_preserves_old_authority_and_bundle(
    tmp_path: Path, monkeypatch
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / "workspace"
    upstream = _filter_authority(workspace)
    argv = _argv(workspace, upstream, _java(tmp_path))
    assert main(argv) == 0
    authority = _authority_path(workspace)
    before = authority.read_bytes()
    manifest = json.loads(before)
    bundle_hashes = {
        role: hashlib.sha256(Path(item["full_path"]).read_bytes()).hexdigest()
        for role, item in manifest["artifacts"].items()
    }
    monkeypatch.setenv("FAIL_POINT", "beagle")
    assert main(argv) == 1
    assert authority.read_bytes() == before
    assert {
        role: hashlib.sha256(Path(item["full_path"]).read_bytes()).hexdigest()
        for role, item in manifest["artifacts"].items()
    } == bundle_hashes


def test_concurrent_success_publishes_one_complete_authority(
    tmp_path: Path, monkeypatch
) -> None:
    from apps.cli.run_wgs_impute import main

    workspace = tmp_path / "workspace"
    argv = _argv(workspace, _filter_authority(workspace), _java(tmp_path))
    barrier = Barrier(2)
    original = __import__(
        "apps.cli.run_wgs_impute", fromlist=["_run_beagle"]
    )._run_beagle

    def synchronized(*args, **kwargs):
        barrier.wait()
        return original(*args, **kwargs)

    monkeypatch.setattr("apps.cli.run_wgs_impute._run_beagle", synchronized)
    results: list[int] = []
    threads = [Thread(target=lambda: results.append(main(argv))) for _ in range(2)]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=20)
    assert sorted(results) in ([0, 0], [0, 1])
    manifest = json.loads(_authority_path(workspace).read_text())
    assert manifest["status"] == "success"
    assert all(Path(item["full_path"]).is_file() for item in manifest["artifacts"].values())