#!/usr/bin/env python3
"""Run Stage 09 SNP/INDEL selection from the authoritative Stage08 VCF."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import fcntl
import gzip
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
from typing import Any
import uuid

from gfa_core.wgs_scientific_identity import (
    build_wgs_analysis_id,
    sha256_file,
)
from gfa_core.wgs_upstream_manifest import resolve_wgs_artifacts


STAGE = "wgs_variant_select"
ROLES = (
    "selected_snp_vcf",
    "selected_snp_vcf_index",
    "selected_indel_vcf",
    "selected_indel_vcf_index",
    "variant_select_report",
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Select SNPs and INDELs from authoritative Stage08 VCF"
    )
    parser.add_argument("--workspace", required=True)
    parser.add_argument("--job-id", required=True)
    parser.add_argument("--prefix", required=True)
    parser.add_argument("--joint-manifest", required=True)
    parser.add_argument(
        "--method", choices=("gatk", "bcftools"), default="gatk"
    )
    parser.add_argument("--gatk", default="gatk")
    parser.add_argument("--bcftools", default="bcftools")
    return parser


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    temporary.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def _safe_component(value: str, label: str) -> None:
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]*", value or ""):
        raise ValueError(f"{label} must be a non-empty safe path component")


def _authority_header(
    path: str | Path,
    *,
    workspace: Path,
    prefix: str,
    job_id: str,
) -> dict[str, Any]:
    try:
        manifest = json.loads(Path(path).read_text(encoding="utf-8"))
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise ValueError(
            "wgs_joint_genotyping authority Manifest is unreadable"
        ) from exc
    if (
        not isinstance(manifest, dict)
        or manifest.get("stage") != "wgs_joint_genotyping"
        or manifest.get("status") != "success"
    ):
        raise ValueError(
            "wgs_joint_genotyping authority Manifest is not successful"
        )
    if (
        manifest.get("prefix") != prefix
        or manifest.get("job_id") != job_id
        or manifest.get("workspace") != str(workspace)
    ):
        raise ValueError(
            "wgs_joint_genotyping authority Manifest identity mismatch"
        )
    return manifest


def _invocation_path(command: str) -> Path:
    resolved = shutil.which(command)
    if not resolved:
        raise RuntimeError(f"executable is unavailable: {command}")
    return Path(resolved).absolute()


def _executable_path(command: str) -> Path:
    return _invocation_path(command).resolve()


def _environment(command: str) -> dict[str, str]:
    return {
        **os.environ,
        "PATH": str(_invocation_path(command).parent)
        + os.pathsep
        + os.environ.get("PATH", ""),
    }


def _checked(
    command: list[str], label: str
) -> subprocess.CompletedProcess[str]:
    result = subprocess.run(
        command,
        text=True,
        capture_output=True,
        check=False,
        env=_environment(command[0]),
    )
    if result.returncode:
        detail = (result.stderr or result.stdout).strip()
        raise RuntimeError(
            f"{label} failed with exit code {result.returncode}: {detail}"
        )
    return result


def _run_logged(command: list[str], log: Any, label: str) -> None:
    log.write(f"{label} argv: {json.dumps(command)}\n")
    log.flush()
    result = subprocess.run(
        command,
        stdout=log,
        stderr=subprocess.STDOUT,
        text=True,
        check=False,
        env=_environment(command[0]),
    )
    log.write(f"{label} exit={result.returncode}\n")
    log.flush()
    if result.returncode:
        raise RuntimeError(f"{label} failed with exit code {result.returncode}")


def _tool_version(command: str, label: str) -> str:
    result = _checked([command, "--version"], f"{label} version check")
    lines = (result.stdout or result.stderr).strip().splitlines()
    if not lines:
        raise RuntimeError(f"{label} version check produced no version")
    return lines[0]


def _vcf_summary(
    path: Path,
    expected_kind: str | None = None,
    *,
    allow_empty: bool = False,
) -> tuple[list[str], str, int, str]:
    if not path.is_file() or path.stat().st_size <= 0:
        raise RuntimeError("VCF is missing or empty")
    samples: list[str] | None = None
    interval = ""
    fallback_interval = ""
    region_row = ""
    records = 0
    try:
        with gzip.open(path, "rt", encoding="utf-8") as handle:
            for raw in handle:
                line = raw.rstrip("\n")
                if line.startswith("#CHROM\t"):
                    fields = line.split("\t")
                    if len(fields) < 10:
                        raise RuntimeError("VCF has no sample columns")
                    samples = fields[9:]
                elif line.startswith("##contig=<ID=") and not region_row:
                    contig = line[len("##contig=<ID="):].split(",", 1)[0]
                    if contig:
                        fallback_interval = f"{contig}:1-1"
                        region_row = f"{contig}\t1\t1\n"
                elif line and not line.startswith("#"):
                    fields = line.split("\t")
                    if len(fields) < 10:
                        raise RuntimeError("VCF record is invalid")
                    ref = fields[3]
                    alts = fields[4].split(",")
                    is_snp = (
                        len(ref) == 1
                        and bool(alts)
                        and all(len(alt) == 1 for alt in alts)
                    )
                    is_indel = any(len(alt) != len(ref) for alt in alts)
                    if expected_kind == "SNP" and not is_snp:
                        raise RuntimeError("selected SNP VCF contains non-SNP")
                    if expected_kind == "INDEL" and not is_indel:
                        raise RuntimeError(
                            "selected INDEL VCF contains non-INDEL"
                        )
                    try:
                        position = int(fields[1])
                    except ValueError as exc:
                        raise RuntimeError(
                            "VCF record position is invalid"
                        ) from exc
                    if not fields[0] or position < 1:
                        raise RuntimeError("VCF record position is invalid")
                    if not interval:
                        interval = f"{fields[0]}:{position}-{position}"
                        region_row = (
                            f"{fields[0]}\t{position}\t{position}\n"
                        )
                    records += 1
    except (OSError, UnicodeError) as exc:
        raise RuntimeError("VCF is unreadable") from exc
    if not interval and allow_empty:
        interval = fallback_interval
    if (
        samples is None
        or not samples
        or len(samples) != len(set(samples))
        or (records < 1 and not allow_empty)
        or not interval
        or not region_row
    ):
        raise RuntimeError("VCF content is incomplete")
    return samples, interval, records, region_row


def _query_index(
    bcftools: str,
    vcf: Path,
    region_row: str,
    scratch: Path,
    log: Any,
    label: str,
) -> None:
    regions = scratch / f".{label}.{uuid.uuid4().hex}.regions.tsv"
    regions.write_text(region_row, encoding="utf-8")
    try:
        _run_logged(
            [bcftools, "view", "-H", "-R", str(regions), str(vcf)],
            log,
            f"{label} indexed query",
        )
    finally:
        regions.unlink(missing_ok=True)


def _select_gatk(
    gatk: str,
    source: Path,
    output: Path,
    kind: str,
    log: Any,
) -> None:
    _run_logged(
        [
            gatk,
            "SelectVariants",
            "-V",
            str(source),
            "-O",
            str(output),
            "-select-type-to-include",
            kind,
        ],
        log,
        f"GATK SelectVariants {kind}",
    )
    _run_logged(
        [gatk, "IndexFeatureFile", "-I", str(output)],
        log,
        f"GATK IndexFeatureFile {kind}",
    )


def _select_bcftools(
    bcftools: str,
    source: Path,
    output: Path,
    kind: str,
    log: Any,
) -> None:
    variant_type = "snps" if kind == "SNP" else "indels"
    _run_logged(
        [
            bcftools,
            "view",
            "-v",
            variant_type,
            "-Oz",
            "-o",
            str(output),
            str(source),
        ],
        log,
        f"bcftools select {kind}",
    )
    _run_logged(
        [bcftools, "index", "--tbi", str(output)],
        log,
        f"bcftools index {kind}",
    )


def _postflight(
    *,
    method: str,
    gatk: str,
    bcftools: str,
    output: Path,
    index: Path,
    kind: str,
    expected_samples: list[str],
    log: Any,
) -> int:
    if index != Path(f"{output}.tbi"):
        raise RuntimeError(f"{kind} index is not its VCF sidecar")
    if not index.is_file() or index.stat().st_size <= 0:
        raise RuntimeError(f"{kind} index is missing or empty")
    samples, interval, records, region_row = _vcf_summary(
        output, kind, allow_empty=True
    )
    if samples != expected_samples:
        raise RuntimeError(f"{kind} sample header mismatch")
    if method == "gatk":
        command = [
            gatk,
            "ValidateVariants",
            "-V",
            str(output),
            "-L",
            interval,
        ]
        command.extend(
            [
                "--validation-type-to-exclude", "ALLELES",
                "--validation-type-to-exclude", "REF",
            ]
        )
        _run_logged(
            command,
            log,
            f"GATK ValidateVariants {kind}",
        )
    else:
        _query_index(
            bcftools, output, region_row, output.parent, log, kind
        )
        _run_logged(
            [bcftools, "stats", str(output)],
            log,
            f"bcftools stats {kind}",
        )
    return records


def _artifact(
    workspace: Path,
    source: Path,
    destination: Path,
    role: str,
) -> dict[str, Any]:
    size = source.stat().st_size
    return {
        "artifact": role,
        "role": role,
        "path": str(destination.relative_to(workspace)),
        "full_path": str(destination.resolve()),
        "type": source.suffix.lstrip(".") or "file",
        "size": size,
        "size_bytes": size,
        "sha256": sha256_file(source),
    }


def _publish(
    bundle: Path,
    destination: Path,
    manifest_work: Path,
    authority: Path,
    started_at: str,
) -> bool:
    lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    with lock_path.open("a+", encoding="utf-8") as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        if authority.is_file():
            try:
                current = json.loads(authority.read_text(encoding="utf-8"))
                if (
                    current.get("status") == "success"
                    and str(current.get("start_time") or "") >= started_at
                ):
                    return False
            except Exception:
                pass
        destination.parent.mkdir(parents=True, exist_ok=True)
        os.replace(bundle, destination)
        os.replace(manifest_work, authority)
    return True


def _failure(
    args: argparse.Namespace, run_id: str, started_at: str, error: str
) -> dict[str, Any]:
    payload = {
        "manifest_version": "1.0",
        "stage": STAGE,
        "status": "failed",
        "prefix": args.prefix,
        "job_id": args.job_id,
        "run_id": run_id,
        "start_time": started_at,
        "method": args.method,
        "artifacts": {},
        "errors": [error],
    }
    if (
        args.method == "gatk"
        and bool(getattr(args, "_fallback_eligible", False))
    ):
        payload["fallback_method"] = "bcftools"
    return payload


def _publish_failure(
    authority: Path,
    args: argparse.Namespace,
    run_id: str,
    started_at: str,
    error: str,
) -> None:
    lock_path = authority.parent.parent / ".locks" / f"{authority.stem}.lock"
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    with lock_path.open("a+", encoding="utf-8") as lock:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX)
        failure_path = (
            authority
            if not authority.exists()
            else authority.parent
            / (
                f"{args.prefix}.{STAGE}.{args.method}.failed."
                f"{run_id}.manifest.json"
            )
        )
        _atomic_json(
            failure_path, _failure(args, run_id, started_at, error)
        )


def run(args: argparse.Namespace) -> int:
    args._fallback_eligible = False
    workspace = Path(args.workspace).resolve()
    run_id = uuid.uuid4().hex
    started_at = datetime.now(timezone.utc).isoformat()
    final = workspace / "final"
    final.mkdir(parents=True, exist_ok=True)
    try:
        _safe_component(args.prefix, "prefix")
    except ValueError as exc:
        _atomic_json(
            final / f"{STAGE}.{args.method}.failed.{run_id}.manifest.json",
            _failure(args, run_id, started_at, str(exc)),
        )
        print(f"{STAGE} failed: {exc}", file=sys.stderr)
        return 1
    authority = (
        final / f"{args.prefix}.{STAGE}.{args.method}.manifest.json"
    )
    try:
        header = _authority_header(
            args.joint_manifest,
            workspace=workspace,
            prefix=args.prefix,
            job_id=args.job_id,
        )
        sample_set_id = str(header.get("wgs_sample_set_id") or "")
        reference_id = str(header.get("wgs_reference_id") or "")
        upstream_analysis_id = str(header.get("wgs_analysis_id") or "")
        if not sample_set_id or not reference_id or not upstream_analysis_id:
            raise ValueError(
                "wgs_joint_genotyping scientific lineage is incomplete"
            )
        resolved = resolve_wgs_artifacts(
            args.joint_manifest,
            "wgs_joint_genotyping",
            ("joint_vcf", "joint_vcf_index"),
            sample_set_id,
            reference_id,
            upstream_analysis_id,
        )
        if any(len(resolved[role]) != 1 for role in resolved):
            raise RuntimeError(
                "authoritative joint VCF and TBI must be singletons"
            )
        joint_vcf = resolved["joint_vcf"][0].path
        joint_index = resolved["joint_vcf_index"][0].path
        if joint_index != Path(f"{joint_vcf}.tbi"):
            raise RuntimeError("joint VCF index is not its declared sidecar")
        expected_samples = (header.get("metrics") or {}).get(
            "output_sample_ids"
        )
        if (
            not isinstance(expected_samples, list)
            or not expected_samples
            or expected_samples != sorted(set(expected_samples))
            or (header.get("metrics") or {}).get("input_sample_ids")
            != expected_samples
        ):
            raise RuntimeError("joint VCF sample set is incomplete")
        (
            observed_samples,
            joint_interval,
            _,
            joint_region_row,
        ) = _vcf_summary(joint_vcf)
        if observed_samples != expected_samples:
            raise RuntimeError("joint VCF sample header mismatch")
        executable = args.gatk if args.method == "gatk" else args.bcftools
        label = "GATK" if args.method == "gatk" else "bcftools"
        try:
            version = _tool_version(executable, label)
        except Exception:
            args._fallback_eligible = args.method == "gatk"
            raise
        parameters = {
            "method": args.method,
            "executable": executable,
            "type_rules": (
                {"SNP": "GATK SNP", "INDEL": "GATK INDEL"}
                if args.method == "gatk"
                else {"SNP": "bcftools snps", "INDEL": "bcftools indels"}
            ),
        }
        analysis_id = build_wgs_analysis_id(
            STAGE,
            [item.sha256 for values in resolved.values() for item in values],
            parameters,
            {args.method: version},
            method=args.method,
        )
        attempt = workspace / ".attempts" / run_id
        bundle = attempt / "bundle"
        bundle.mkdir(parents=True, exist_ok=False)
        base = f"{args.prefix}.{STAGE}.cohort.selected"
        snp = bundle / f"{base}.snp.vcf.gz"
        indel = bundle / f"{base}.indel.vcf.gz"
        snp_index = Path(f"{snp}.tbi")
        indel_index = Path(f"{indel}.tbi")
        report = bundle / f"{args.prefix}.{STAGE}.cohort.report.md"
        log_path = attempt / "variant-select.log"
        with log_path.open("w", encoding="utf-8") as log:
            args._fallback_eligible = args.method == "gatk"
            if args.method == "gatk":
                _run_logged(
                    [
                        args.gatk,
                        "ValidateVariants",
                        "-V",
                        str(joint_vcf),
                        "-L",
                        joint_interval,
                        "--validation-type-to-exclude",
                        "ALLELES",
                        "--validation-type-to-exclude",
                        "REF",
                    ],
                    log,
                    "GATK ValidateVariants input",
                )
            else:
                _query_index(
                    args.bcftools,
                    joint_vcf,
                    joint_region_row,
                    attempt,
                    log,
                    "input",
                )
            for output, kind in ((snp, "SNP"), (indel, "INDEL")):
                if args.method == "gatk":
                    _select_gatk(
                        args.gatk, joint_vcf, output, kind, log
                    )
                else:
                    _select_bcftools(
                        args.bcftools, joint_vcf, output, kind, log
                    )
            snp_records = _postflight(
                method=args.method,
                gatk=args.gatk,
                bcftools=args.bcftools,
                output=snp,
                index=snp_index,
                kind="SNP",
                expected_samples=expected_samples,
                log=log,
            )
            indel_records = _postflight(
                method=args.method,
                gatk=args.gatk,
                bcftools=args.bcftools,
                output=indel,
                index=indel_index,
                kind="INDEL",
                expected_samples=expected_samples,
                log=log,
            )
        args._fallback_eligible = False
        report.write_text(
            "# WGS variant selection\n\n"
            f"- method: {args.method}\n"
            f"- SNP records: {snp_records}\n"
            f"- INDEL records: {indel_records}\n"
            "- Methods are recorded as executed; scientific equivalence is "
            "not asserted.\n",
            encoding="utf-8",
        )
        destination = (
            final
            / ".wgs_variant_select_runs"
            / args.prefix
            / args.method
            / run_id
        )
        sources = {
            "selected_snp_vcf": snp,
            "selected_snp_vcf_index": snp_index,
            "selected_indel_vcf": indel,
            "selected_indel_vcf_index": indel_index,
            "variant_select_report": report,
        }
        artifacts = {
            role: _artifact(
                workspace,
                source,
                destination / source.name,
                role,
            )
            for role, source in sources.items()
        }
        manifest = {
            "manifest_version": "1.0",
            "stage": STAGE,
            "status": "success",
            "prefix": args.prefix,
            "job_id": args.job_id,
            "run_id": run_id,
            "start_time": started_at,
            "workspace": str(workspace),
            "method": args.method,
            "wgs_sample_set_id": sample_set_id,
            "wgs_reference_id": reference_id,
            "upstream_wgs_analysis_ids": {
                "wgs_joint_genotyping": upstream_analysis_id
            },
            "wgs_analysis_id": analysis_id,
            "parameters": parameters,
            "software": {
                args.method: {
                    "path": str(_executable_path(executable)),
                    "version": version,
                }
            },
            "metrics": {
                "input_sample_ids": expected_samples,
                "output_sample_ids": expected_samples,
                "total_samples": len(expected_samples),
                "snp_records": snp_records,
                "indel_records": indel_records,
            },
            "artifacts": artifacts,
        }
        manifest_work = attempt / "manifest.json"
        _atomic_json(manifest_work, manifest)
        if not _publish(
            bundle, destination, manifest_work, authority, started_at
        ):
            raise RuntimeError(
                "stale variant-select attempt cannot replace newer method run"
            )
        return 0
    except Exception as exc:
        _publish_failure(
            authority, args, run_id, started_at, str(exc)
        )
        print(f"{STAGE} failed: {exc}", file=sys.stderr)
        return 1


def main(argv: list[str] | None = None) -> int:
    return run(build_parser().parse_args(argv))


if __name__ == "__main__":
    raise SystemExit(main())