#!/usr/bin/env python3
"""Run Stage 02 WGS reference indexing with samtools, GATK and BWA-MEM2."""

from __future__ import annotations

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

from gfa_core.wgs_scientific_identity import build_wgs_analysis_id, build_wgs_reference_id, sha256_file


STAGE = "wgs_index"
COMPONENTS = (".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Build the required WGS reference indices")
    parser.add_argument("--workspace", required=True)
    parser.add_argument("--job-id", required=True)
    parser.add_argument("--prefix", required=True)
    parser.add_argument("--reference-fasta", required=True)
    parser.add_argument("--samtools", default="samtools")
    parser.add_argument("--bwa-mem2", default="bwa-mem2")
    parser.add_argument("--gatk", default="gatk")
    parser.add_argument("--threads", type=int, default=1)
    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 _authority_lock_path(manifest_path: Path) -> Path:
    return manifest_path.parent.parent / ".locks" / f"{manifest_path.stem}.lock"


def _is_safe_path_component(value: Any) -> bool:
    return (
        isinstance(value, str)
        and bool(value)
        and value not in {".", ".."}
        and "/" not in value
        and "\\" not in value
        and not any(ord(character) < 32 or ord(character) == 127 for character in value)
    )


def _is_generated_run_id(value: Any) -> bool:
    return (
        _is_safe_path_component(value)
        and len(value) == 32
        and all(character in "0123456789abcdef" for character in value)
    )


def _validate_prefix(prefix: Any) -> None:
    if not _is_safe_path_component(prefix):
        raise ValueError("prefix must be a non-empty safe path component")


def _publish_transaction(bundle_work: Path, bundle_destination: Path, manifest_work: Path, manifest_path: Path, started_at: str) -> bool:
    """Publish a complete run bundle before switching the authority Manifest."""
    lock_path = _authority_lock_path(manifest_path)
    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 manifest_path.is_file():
            try:
                existing = json.loads(manifest_path.read_text(encoding="utf-8"))
                if existing.get("status") == "success" and str(existing.get("start_time") or "") > started_at:
                    return False
            except Exception:
                pass
        bundle_destination.parent.mkdir(parents=True, exist_ok=True)
        os.replace(bundle_work, bundle_destination)
        os.replace(manifest_work, manifest_path)
    return True


def _tool_version(tool: str, *, bwa_mem2: bool = False) -> str:
    result = subprocess.run([tool, "version" if bwa_mem2 else "--version"], text=True, capture_output=True, check=False)
    if result.returncode:
        raise RuntimeError(f"{tool} version check failed with exit code {result.returncode}: {result.stderr.strip()}")
    lines = (result.stdout or result.stderr).strip().splitlines()
    if not lines:
        raise RuntimeError(f"{tool} version check produced no version")
    return lines[0]


def _validate_reference(path: Path) -> list[tuple[str, int]]:
    if not path.is_file() or path.stat().st_size <= 0:
        raise ValueError("reference FASTA must be an existing non-empty regular file")
    seen: set[str] = set()
    contigs: list[tuple[str, int]] = []
    current = ""
    length = 0
    with path.open("r", encoding="utf-8") as handle:
        for line_number, line in enumerate(handle, 1):
            text = line.strip()
            if not text:
                continue
            if text.startswith(">"):
                if current:
                    if length <= 0:
                        raise ValueError(f"reference FASTA contig {current!r} has no sequence")
                    contigs.append((current, length))
                current = text[1:].split(maxsplit=1)[0]
                if not current:
                    raise ValueError(f"reference FASTA header at line {line_number} has no contig name")
                if current in seen:
                    raise ValueError(f"reference FASTA contig names must be unique: {current}")
                seen.add(current)
                length = 0
            else:
                if not current:
                    raise ValueError(f"reference FASTA sequence at line {line_number} precedes a header")
                length += len("".join(text.split()))
    if not current or length <= 0:
        raise ValueError("reference FASTA must contain at least one header with sequence")
    contigs.append((current, length))
    return contigs


def _validate_fai(path: Path, expected: list[tuple[str, int]]) -> None:
    observed: list[tuple[str, int]] = []
    try:
        with path.open("r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, 1):
                fields = line.rstrip("\n").split("\t")
                if len(fields) < 5 or not fields[0]:
                    raise ValueError(f"invalid FAI record at line {line_number}")
                observed.append((fields[0], int(fields[1])))
    except (OSError, UnicodeError, ValueError) as exc:
        raise RuntimeError(f"FAI content validation failed: {exc}") from exc
    if observed != expected:
        raise RuntimeError(f"FAI contig/length mismatch: expected {expected}, observed {observed}")


def _validate_sequence_dictionary(path: Path, expected: list[tuple[str, int]]) -> None:
    observed: list[tuple[str, int]] = []
    try:
        with path.open("r", encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, 1):
                if not line.startswith("@SQ\t"):
                    continue
                fields = line.rstrip("\n").split("\t")[1:]
                names = [field[3:] for field in fields if field.startswith("SN:")]
                lengths = [field[3:] for field in fields if field.startswith("LN:")]
                if len(names) != 1 or len(lengths) != 1:
                    raise ValueError(f"invalid dictionary @SQ record at line {line_number}")
                observed.append((names[0], int(lengths[0])))
    except (OSError, UnicodeError, ValueError) as exc:
        raise RuntimeError(f"sequence dictionary content validation failed: {exc}") from exc
    if observed != expected:
        raise RuntimeError(f"sequence dictionary contig/length mismatch: expected {expected}, observed {observed}")


def _run_tool(command: list[str], name: str) -> None:
    result = subprocess.run(command, text=True, capture_output=True, check=False)
    if result.returncode:
        raise RuntimeError(f"{name} failed with exit code {result.returncode}: {result.stderr.strip()}")


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


def _required_paths(reference: Path, dictionary: Path) -> dict[str, Path]:
    output = {"reference_fai": Path(f"{reference}.fai"), "reference_dict": dictionary}
    output.update({f"bwa_mem2_index_component.{component}": Path(f"{reference}{component}") for component in COMPONENTS})
    return output


def _has_symlink_component(path: Path, boundary: Path) -> bool:
    try:
        relative = path.relative_to(boundary)
    except ValueError:
        return True
    current = boundary
    for part in relative.parts:
        current = current / part
        if current.is_symlink():
            return True
    return False


def _valid_reuse(manifest_path: Path, *, workspace: Path, job_id: str, prefix: str, reference_id: str, analysis_id: str, parameters: dict[str, Any]) -> bool:
    try:
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    except Exception:
        return False
    run_id = manifest.get("run_id")
    if not all((manifest.get("stage") == STAGE, manifest.get("status") == "success", manifest.get("job_id") == job_id,
                manifest.get("prefix") == prefix, manifest.get("wgs_reference_id") == reference_id,
                manifest.get("wgs_analysis_id") == analysis_id, manifest.get("parameters") == parameters,
                _is_safe_path_component(prefix), _is_generated_run_id(run_id))):
        return False
    workspace = workspace.resolve()
    bundle_root = workspace / "final" / ".wgs_index_runs" / prefix / run_id
    if _has_symlink_component(bundle_root, workspace):
        return False
    bundle_root_resolved = bundle_root.resolve()
    artifacts = manifest.get("artifacts") if isinstance(manifest.get("artifacts"), dict) else {}
    expected = {"reference_fasta", "reference_fai", "reference_dict", *{f"bwa_mem2_index_component.{component}" for component in COMPONENTS}}
    if set(artifacts) != expected:
        return False
    for key, item in artifacts.items():
        if not isinstance(item, dict):
            return False
        expected_role = "bwa_mem2_index_component" if key.startswith("bwa_mem2_index_component.") else key
        if str(item.get("role") or item.get("artifact") or "") != expected_role:
            return False
        expected_component = key.removeprefix("bwa_mem2_index_component.") if expected_role == "bwa_mem2_index_component" else ""
        if str(item.get("component") or "") != expected_component:
            return False
        relative_path = Path(str(item.get("path") or ""))
        full_path = Path(str(item.get("full_path") or ""))
        if not str(relative_path) or relative_path.is_absolute() or ".." in relative_path.parts or not full_path.is_absolute() or ".." in full_path.parts:
            return False
        path_from_relative = workspace / relative_path
        if _has_symlink_component(path_from_relative, workspace) or _has_symlink_component(full_path, workspace):
            return False
        resolved_path = path_from_relative.resolve()
        resolved_full_path = full_path.resolve()
        if resolved_path != resolved_full_path:
            return False
        try:
            resolved_path.relative_to(bundle_root_resolved)
        except ValueError:
            return False
        if not resolved_path.is_file() or resolved_path.stat().st_size <= 0:
            return False
        actual_size = resolved_path.stat().st_size
        if item.get("size") != actual_size or item.get("size_bytes") != actual_size:
            return False
        declared_hash = str(item.get("sha256") or "").lower()
        if len(declared_hash) != 64 or any(char not in "0123456789abcdef" for char in declared_hash) or sha256_file(resolved_path) != declared_hash:
            return False
    return True


def _refresh_reuse_authority(
    manifest_path: Path, *, workspace: Path, job_id: str, prefix: str, reference_id: str,
    analysis_id: str, parameters: dict[str, Any], reused_at: str, reuse_attempt_id: str,
) -> bool:
    lock_path = _authority_lock_path(manifest_path)
    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 not _valid_reuse(
            manifest_path, workspace=workspace, job_id=job_id, prefix=prefix,
            reference_id=reference_id, analysis_id=analysis_id, parameters=parameters,
        ):
            return False
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        previous = manifest.get("reuse") if isinstance(manifest.get("reuse"), dict) else {}
        manifest["reuse"] = {
            "count": int(previous.get("count") or 0) + 1,
            "last_reused_at": reused_at,
            "last_attempt_id": reuse_attempt_id,
        }
        _atomic_json(manifest_path, manifest)
    return True


def _failure_manifest(args: argparse.Namespace, error: str, run_id: str, started_at: str) -> dict[str, Any]:
    return {"manifest_version": "1.0", "stage": STAGE, "status": "failed", "prefix": args.prefix,
            "job_id": args.job_id, "run_id": run_id, "start_time": started_at, "artifacts": {}, "errors": [error]}


def run(args: argparse.Namespace) -> int:
    workspace = Path(args.workspace).resolve()
    try:
        _validate_prefix(args.prefix)
    except ValueError as exc:
        print(f"wgs_index failed: {exc}", file=sys.stderr)
        return 1
    final_dir = workspace / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = final_dir / f"{args.prefix}.{STAGE}.manifest.json"
    run_id, started_at = uuid.uuid4().hex, datetime.now(timezone.utc).isoformat()
    try:
        if args.threads < 1:
            raise ValueError("threads must be positive")
        reference_input = Path(args.reference_fasta).resolve()
        reference_contigs = _validate_reference(reference_input)
        reference_sha256 = sha256_file(reference_input)
        reference_id = build_wgs_reference_id(reference_sha256)
        versions = {"samtools": _tool_version(args.samtools), "bwa_mem2": _tool_version(args.bwa_mem2, bwa_mem2=True), "gatk": _tool_version(args.gatk)}
        parameters = {"samtools": args.samtools, "bwa_mem2": args.bwa_mem2, "gatk": args.gatk, "threads": args.threads, "components": list(COMPONENTS)}
        analysis_id = build_wgs_analysis_id(STAGE, [reference_sha256], parameters, versions, method="samtools_gatk_bwa_mem2")
        if _refresh_reuse_authority(
            manifest_path, workspace=workspace, job_id=args.job_id, prefix=args.prefix,
            reference_id=reference_id, analysis_id=analysis_id, parameters=parameters,
            reused_at=started_at, reuse_attempt_id=run_id,
        ):
            return 0

        attempt = workspace / ".attempts" / run_id
        bundle_work = attempt / "bundle"
        bundle_work.mkdir(parents=True, exist_ok=False)
        reference = bundle_work / f"{args.prefix}.{STAGE}.reference.fasta"
        shutil.copy2(reference_input, reference)
        dictionary = bundle_work / f"{args.prefix}.{STAGE}.reference.dict"
        _run_tool([args.samtools, "faidx", str(reference)], "samtools faidx")
        _run_tool([args.gatk, "CreateSequenceDictionary", "-R", str(reference), "-O", str(dictionary)], "GATK CreateSequenceDictionary")
        _run_tool([args.bwa_mem2, "index", str(reference)], "bwa-mem2 index")
        outputs = _required_paths(reference, dictionary)
        missing = [key for key, path in outputs.items() if not path.is_file() or path.stat().st_size <= 0]
        if missing:
            raise RuntimeError(f"missing or empty required index output(s): {', '.join(missing)}")
        _validate_fai(outputs["reference_fai"], reference_contigs)
        _validate_sequence_dictionary(outputs["reference_dict"], reference_contigs)
        destination = final_dir / ".wgs_index_runs" / args.prefix / run_id
        artifacts = {"reference_fasta": _artifact(workspace, reference, destination / reference.name, "reference_fasta")}
        for key, source in outputs.items():
            role = "bwa_mem2_index_component" if key.startswith("bwa_mem2_index_component.") else key
            component = key.removeprefix("bwa_mem2_index_component.") if role == "bwa_mem2_index_component" else ""
            artifacts[key] = _artifact(workspace, source, destination / source.name, role, component)
        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), "wgs_reference_id": reference_id,
                    "wgs_analysis_id": analysis_id, "parameters": parameters,
                    "software": {name: {"path": getattr(args, name if name != "bwa_mem2" else "bwa_mem2"), "version": version} for name, version in versions.items()},
                    "metrics": {"required_index_outputs": 7, "components": list(COMPONENTS)}, "artifacts": artifacts}
        manifest_work = attempt / "manifest.json"
        _atomic_json(manifest_work, manifest)
        if not _publish_transaction(bundle_work, destination, manifest_work, manifest_path, started_at):
            raise RuntimeError("stale WGS index attempt cannot replace a newer successful run")
        return 0
    except Exception as exc:
        failure = _failure_manifest(args, str(exc), run_id, started_at)
        failure_path = manifest_path if not manifest_path.exists() else final_dir / f"{args.prefix}.{STAGE}.failed.{run_id}.manifest.json"
        _atomic_json(failure_path, failure)
        print(f"wgs_index 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())