#!/usr/bin/env python3
"""Calculate and validate a TASSEL centered-IBS kinship matrix."""

from __future__ import annotations

import argparse
import gzip
import hashlib
import json
import os
import platform
import re
import shlex
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np

from apps.cli.run_gwas_tassel import _resolved_command, _tassel_version
from gfa_core.gwas_scientific_identity import build_identity_recipe, materialize_stage_identity, new_run_id


STAGE = "gwas_kinship"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def resolve_executable(value: str) -> str:
    return _resolved_command(value)


def tassel_version(value: str) -> str:
    return _tassel_version(value)


def _write_json(path: Path, value: dict[str, Any]) -> None:
    path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def _vcf_samples(path: Path) -> list[str]:
    opener = gzip.open if path.suffix == ".gz" else open
    with opener(path, "rt", encoding="utf-8") as handle:
        for line in handle:
            if line.startswith("#CHROM\t"):
                samples = line.rstrip("\n").split("\t")[9:]
                if not samples or len(samples) != len(set(samples)):
                    raise ValueError("sorted VCF sample IDs must be nonempty and unique")
                return samples
    raise ValueError("sorted VCF #CHROM header is missing")


def validate_kinship_matrix(path: Path, expected_samples: list[str]) -> dict[str, Any]:
    lines = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    body = [line for line in lines if not line.startswith("##")]
    if not body:
        raise ValueError("kinship matrix is empty")
    try:
        declared_count = int(body[0])
    except ValueError as exc:
        raise ValueError("kinship matrix sample count is invalid") from exc
    rows = body[1:]
    if declared_count != len(rows) or declared_count != len(expected_samples):
        raise ValueError("kinship matrix is not square for the expected sample count")
    ids: list[str] = []
    values: list[list[float]] = []
    for line in rows:
        fields = line.split()
        if len(fields) != declared_count + 1:
            raise ValueError("kinship matrix row width is invalid")
        ids.append(fields[0])
        try:
            values.append([float(value) for value in fields[1:]])
        except ValueError as exc:
            raise ValueError("kinship matrix contains nonnumeric values") from exc
    matrix = np.asarray(values, dtype=float)
    if not np.isfinite(matrix).all():
        raise ValueError("kinship matrix contains nonfinite values")
    if ids != expected_samples:
        raise ValueError("kinship matrix sample order does not match sorted VCF")
    if not np.allclose(matrix, matrix.T, rtol=1e-7, atol=1e-7):
        raise ValueError("kinship matrix must be symmetric")
    minimum_eigenvalue = float(np.linalg.eigvalsh((matrix + matrix.T) / 2.0).min())
    if minimum_eigenvalue < -1e-6:
        raise ValueError("kinship matrix is not positive-semidefinite")
    return {
        "sample_count": declared_count,
        "sample_order_conservation": True,
        "symmetric": True,
        "positive_semidefinite": True,
        "minimum_eigenvalue": minimum_eigenvalue,
        "matrix_type": "Centered_IBS",
    }


def run_tassel_kinship(command: list[str], output: Path) -> str:
    completed = subprocess.run(command, text=True, capture_output=True, check=False)
    combined = (completed.stdout or "") + (completed.stderr or "")
    if completed.returncode != 0:
        raise RuntimeError(f"TASSEL kinship failed with exit code {completed.returncode}")
    if not output.is_file() or output.stat().st_size == 0:
        raise RuntimeError("TASSEL completed without a kinship matrix")
    return combined


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


def _runner_command(args: argparse.Namespace, workspace: Path, vcf: Path) -> list[str]:
    command = [
        "python", "-m", "apps.cli.run_gwas_kinship",
        "--workspace", str(workspace), "--job-id", args.job_id,
        "--prefix", args.prefix, "--vcf", str(vcf),
        "--tassel-pipeline", args.tassel_pipeline, "--xmx", args.xmx,
    ]
    if args.genotype_scope_id:
        command.extend(["--genotype-scope-id", args.genotype_scope_id])
    if args.log_file:
        command.extend(["--log-file", str(args.log_file)])
    return command


def _write_reproducibility(
    workspace: Path,
    identity12: str,
    run_id: str,
    command: list[str],
    manifest_path: Path,
    manifest: dict[str, Any],
    environment_text: str,
) -> None:
    bundle = workspace / "reproducibility" / STAGE / f"{manifest['prefix']}.{identity12}" / run_id
    bundle.mkdir(parents=True, exist_ok=True)
    canonical = shlex.join(command)
    (bundle / "commands.sh").write_text("#!/usr/bin/env bash\nset -euo pipefail\n" + canonical + "\n", encoding="utf-8")
    (bundle / "environment.txt").write_text(environment_text, encoding="utf-8")
    _write_json(bundle / "progress.json", {
        "stage": STAGE, "status": manifest["status"], "run_id": run_id,
        "manifest_path": str(manifest_path.relative_to(workspace)),
    })
    _write_json(bundle / "provenance.json", {
        "stage": STAGE, "run_id": run_id, "runner_command": canonical,
        "software": manifest.get("software", {}), "artifacts": manifest.get("artifacts", {}),
    })
    targets = [manifest_path]
    for artifact in manifest.get("artifacts", {}).values():
        candidate = workspace / artifact["path"]
        if candidate.is_file():
            targets.append(candidate)
    (bundle / "checksums.sha256").write_text(
        "".join(f"{sha256_file(path)}  {path.relative_to(workspace)}\n" for path in sorted(targets)),
        encoding="utf-8",
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--workspace", required=True, type=Path)
    parser.add_argument("--job-id", required=True)
    parser.add_argument("--prefix", required=True)
    parser.add_argument("--vcf", required=True, type=Path)
    parser.add_argument("--genotype-scope-id", default="")
    parser.add_argument("--tassel-pipeline", default="run_pipeline.pl")
    parser.add_argument("--xmx", default="10g")
    parser.add_argument("--log-file", type=Path)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    workspace = args.workspace.resolve()
    final_dir = workspace / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    vcf = args.vcf.resolve()
    run_id = new_run_id(STAGE)
    started = datetime.now(timezone.utc)
    started_clock = time.monotonic()
    failed_manifest = final_dir / f"{args.prefix}.{STAGE}.failed.manifest.json"
    command = _runner_command(args, workspace, vcf)
    manifest: dict[str, Any] = {
        "manifest_version": "1.0", "stage": STAGE, "prefix": args.prefix,
        "job_id": args.job_id, "run_id": run_id, "status": "failed",
        "timestamp": started.isoformat(),
        "start_time": started.isoformat(), "end_time": started.isoformat(),
        "elapsed_sec": 0.0, "artifacts": {}, "cmd": command,
        "workspace_root": str(workspace), "workspace": str(final_dir),
        "input_dir": str(workspace / "input"),
    }
    environment_text = ""
    try:
        if not vcf.is_file():
            raise FileNotFoundError("Authoritative sorted VCF does not exist")
        vcf_sha = sha256_file(vcf)
        genotype_scope = args.genotype_scope_id or f"genotype:{vcf_sha}"
        if not re.fullmatch(r"genotype:[0-9a-f]{64}", genotype_scope):
            raise ValueError("Provided genotype scope is invalid")
        samples = _vcf_samples(vcf)
        executable = resolve_executable(args.tassel_pipeline)
        if not Path(executable).is_file() or not os.access(executable, os.X_OK):
            raise FileNotFoundError("TASSEL pipeline executable was not found or is not executable")
        version = tassel_version(executable)
        if version == "unknown":
            raise RuntimeError("TASSEL version could not be proven")
        recipe = build_identity_recipe(
            stage=STAGE, required_artifact_roles=["sorted_vcf"],
            parameters={"kinship_method": "Centered_IBS", "xmx": args.xmx},
            method_version="tassel_centered_ibs.v1", contract_version="v1",
            tool_versions={"tassel": version},
        )
        identity = materialize_stage_identity(
            recipe, {"sorted_vcf": {"sha256": vcf_sha}},
            {"genotype_scope_id": genotype_scope},
        )
        identity12 = identity["analysis_scope_id"].split(":", 1)[1][:12]
        attempt_dir = workspace / ".attempts" / run_id
        attempt_dir.mkdir(parents=True, exist_ok=False)
        matrix_work = attempt_dir / "kinship.txt"
        tassel_command = [
            executable, f"-Xmx{args.xmx}", "-fork1", "-vcf", str(vcf),
            "-ck", "-export", str(matrix_work), "-runfork1",
        ]
        run_tassel_kinship(tassel_command, matrix_work)
        metrics = validate_kinship_matrix(matrix_work, samples)

        stem = f"{args.prefix}.{STAGE}.{identity12}"
        matrix_final = final_dir / f"{stem}.kinship.txt"
        environment_final = final_dir / f"{stem}.environment.txt"
        report_final = final_dir / f"{stem}.report.md"
        manifest_path = final_dir / f"{stem}.manifest.json"
        shutil.copy2(matrix_work, matrix_final)
        environment_text = "\n".join([
            "runner=apps.cli.run_gwas_kinship.py",
            f"runtime_environment={os.environ.get('CONDA_DEFAULT_ENV', 'unknown')}",
            f"python_resolved={Path(sys.executable).resolve()}",
            f"python_version={platform.python_version()}",
            f"tassel_resolved={executable}", f"tassel_version={version}",
            f"platform={platform.platform()}",
        ]) + "\n"
        environment_final.write_text(environment_text, encoding="utf-8")
        report_final.write_text(
            "# GWAS TASSEL kinship\n\n"
            f"- Genotype scope: `{genotype_scope}`\n"
            f"- Analysis scope: `{identity['analysis_scope_id']}`\n"
            f"- Samples: {metrics['sample_count']}\n"
            f"- Matrix type: `Centered_IBS`\n"
            f"- Minimum eigenvalue: {metrics['minimum_eigenvalue']:.10g}\n",
            encoding="utf-8",
        )
        manifest.update({
            "status": "success", "genotype_scope_id": genotype_scope,
            "analysis_scope_id": identity["analysis_scope_id"],
            "identity_recipe": recipe,
            "inputs": {"sorted_vcf": {"path": str(vcf), "size": vcf.stat().st_size, "sha256": vcf_sha}},
            "parameters": {"kinship_method": "Centered_IBS", "xmx": args.xmx},
            "metrics": metrics,
            "software": {"tassel": {"resolved_path": executable, "version": version}},
            "artifacts": {
                "kinship_matrix": _artifact(workspace, matrix_final, "kinship_matrix", "text/tab-separated-values"),
                "environment": _artifact(workspace, environment_final, "environment", "text/plain"),
                "report": _artifact(workspace, report_final, "report", "text/markdown"),
            },
        })
        manifest["end_time"] = datetime.now(timezone.utc).isoformat()
        manifest["elapsed_sec"] = round(time.monotonic() - started_clock, 6)
        _write_json(manifest_path, manifest)
        failed_manifest.unlink(missing_ok=True)
        _write_reproducibility(workspace, identity12, run_id, command, manifest_path, manifest, environment_text)
        return 0
    except Exception as exc:
        manifest["error"] = str(exc)
        manifest["end_time"] = datetime.now(timezone.utc).isoformat()
        manifest["elapsed_sec"] = round(time.monotonic() - started_clock, 6)
        _write_json(failed_manifest, manifest)
        if args.log_file:
            args.log_file.parent.mkdir(parents=True, exist_ok=True)
            args.log_file.write_text(f"ERROR: {exc}\n", encoding="utf-8")
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1


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