#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import re
import subprocess
import sys
import tarfile
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


SAFE_ACCESSION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
REQUIRED_COLUMNS = ("accession", "gff3", "pep", "cds")


def resolve_input(value: str, samples_path: Path) -> Path:
    path = Path(value)
    candidates = [path] if path.is_absolute() else [samples_path.parent / path, Path.cwd() / path]
    for candidate in candidates:
        if candidate.is_file():
            return candidate.resolve()
    raise FileNotFoundError(f"input file does not exist: {value}")


def load_samples(samples_path: Path) -> list[dict[str, str | Path]]:
    with samples_path.open(encoding="utf-8-sig", newline="") as handle:
        reader = csv.DictReader(handle, delimiter="\t")
        missing = [column for column in REQUIRED_COLUMNS if column not in (reader.fieldnames or [])]
        if missing:
            raise ValueError(f"samples.tsv is missing columns: {', '.join(missing)}")
        rows = list(reader)
    if not rows:
        raise ValueError("samples.tsv contains no samples")
    seen: set[str] = set()
    normalized: list[dict[str, str | Path]] = []
    for row in rows:
        accession = (row.get("accession") or "").strip()
        if not SAFE_ACCESSION.fullmatch(accession):
            raise ValueError(f"unsafe accession token: {accession!r}")
        if accession in seen:
            raise ValueError(f"duplicate accession: {accession}")
        seen.add(accession)
        normalized.append({
            "accession": accession,
            "gff3": resolve_input((row.get("gff3") or "").strip(), samples_path),
            "pep": resolve_input((row.get("pep") or "").strip(), samples_path),
            "cds": resolve_input((row.get("cds") or "").strip(), samples_path),
        })
    return normalized


def run_one(row: dict[str, str | Path], script: Path, output_dir: Path) -> dict[str, str | int]:
    accession = str(row["accession"])
    prefix = output_dir / accession
    command = [
        sys.executable, str(script),
        "--accession", accession,
        "--gff3", str(row["gff3"]),
        "--pep", str(row["pep"]),
        "--cds", str(row["cds"]),
        "--out-prefix", str(prefix),
    ]
    completed = subprocess.run(command, text=True, capture_output=True)
    if completed.returncode != 0:
        raise RuntimeError(f"{accession} failed:\n{completed.stderr or completed.stdout}")
    outputs = [
        Path(f"{prefix}.primary.pep.fa"),
        Path(f"{prefix}.primary.cds.fa"),
        Path(f"{prefix}.id_mapping.tsv"),
    ]
    for output in outputs:
        if not output.is_file():
            raise RuntimeError(f"{accession} did not create {output.name}")
    return {
        "accession": accession,
        "primary_pep": outputs[0].name,
        "primary_cds": outputs[1].name,
        "id_mapping": outputs[2].name,
        "primary_pep_bytes": outputs[0].stat().st_size,
        "primary_cds_bytes": outputs[1].stat().st_size,
    }


def write_manifest(rows: list[dict[str, str | int]], path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = ["accession", "primary_pep", "primary_cds", "id_mapping", "primary_pep_bytes", "primary_cds_bytes"]
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields, delimiter="\t", lineterminator="\n")
        writer.writeheader()
        writer.writerows(sorted(rows, key=lambda row: str(row["accession"])))


def main() -> None:
    parser = argparse.ArgumentParser(description="Extract primary isoforms for every row in samples.tsv")
    parser.add_argument("--samples", required=True, type=Path)
    parser.add_argument("--jobs", type=int, default=4, help="maximum concurrent samples; does not limit sample count")
    parser.add_argument("--output-prefix", required=True, type=Path)
    args = parser.parse_args()
    if args.jobs < 1:
        parser.error("--jobs must be at least 1")
    samples_path = args.samples.resolve()
    rows = load_samples(samples_path)
    script = Path(__file__).with_name("select_primary_isoform.py").resolve()
    archive_path = Path(f"{args.output_prefix}.primary.tar.gz")
    manifest_path = Path(f"{args.output_prefix}.primary_manifest.tsv")
    archive_path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(prefix="primary_batch_") as temp_name:
        output_dir = Path(temp_name) / "primary"
        output_dir.mkdir()
        results: list[dict[str, str | int]] = []
        with ThreadPoolExecutor(max_workers=min(args.jobs, len(rows))) as executor:
            futures = {executor.submit(run_one, row, script, output_dir): str(row["accession"]) for row in rows}
            for future in as_completed(futures):
                result = future.result()
                results.append(result)
                print(f"completed primary isoforms: {result['accession']}", flush=True)
        write_manifest(results, manifest_path)
        with tarfile.open(archive_path, "w:gz") as archive:
            for file_path in sorted(output_dir.iterdir()):
                archive.add(file_path, arcname=f"primary/{file_path.name}")
    print(f"processed {len(rows)} samples")
    print(f"archive: {archive_path}")
    print(f"manifest: {manifest_path}")


if __name__ == "__main__":
    main()