# -*- coding: utf-8 -*-
"""
JanusX: FASTQ -> VCF Pipeline Command-Line Interface

This CLI prepares metadata from a FASTQ directory and runs:
  fastp -> bwa-mem -> markdup -> gVCF -> joint-calling -> SNP filtering -> TSV

Requirements
------------
- Paired-end FASTQ files named as: *.R1.fastq.gz / *.R2.fastq.gz
- Reference FASTA with a .fai index (auto-generated by this CLI)

Examples
--------
  jx fastq2vcf \\
    --reference /data/ref.fa \\
    --fastq-dir /data/fastq \\
    --workdir /data/out \\
    --backend nohup \\
    --maxtask 2
"""

import os
import re
import time
import socket
import argparse
import sys
import subprocess
import shutil
import urllib.request
from pathlib import Path
from collections import defaultdict
from typing import List, Optional

import janusx.pipeline
from janusx.pipeline.fastq2vcf import fastq2vcf, indexREF
from janusx.pipeline.pipeline import wrap_cmd
from ._common.log import setup_logging
from ._common.config_render import emit_cli_configuration
from ._common.pathcheck import ensure_dir_exists, ensure_file_exists
from ._common.status import (
    CliStatus,
    get_rich_spinner_name,
    print_success,
    print_failure,
    format_elapsed,
)

try:
    from rich.progress import (
        Progress,
        SpinnerColumn,
        TextColumn,
        BarColumn,
        DownloadColumn,
        TransferSpeedColumn,
        TimeElapsedColumn,
        TimeRemainingColumn,
    )
    _HAS_RICH_PROGRESS = True
except Exception:
    Progress = None  # type: ignore[assignment]
    SpinnerColumn = None  # type: ignore[assignment]
    TextColumn = None  # type: ignore[assignment]
    BarColumn = None  # type: ignore[assignment]
    DownloadColumn = None  # type: ignore[assignment]
    TransferSpeedColumn = None  # type: ignore[assignment]
    TimeElapsedColumn = None  # type: ignore[assignment]
    TimeRemainingColumn = None  # type: ignore[assignment]
    _HAS_RICH_PROGRESS = False

try:
    from tqdm.auto import tqdm
    _HAS_TQDM = True
except Exception:
    tqdm = None  # type: ignore[assignment]
    _HAS_TQDM = False

READ_RE = re.compile(r"\.(R[12])\.(fastq|fq)(\.gz)?$", re.IGNORECASE)
FASTQ_SUFFIXES = (".fastq", ".fq", ".fastq.gz", ".fq.gz")


class _DownloadProgress:
    """
    Download progress adapter: rich first, tqdm fallback.
    """
    def __init__(self, total: Optional[int], desc: str = "Downloading") -> None:
        self.total = total
        self.desc = desc
        self._backend = "none"
        self._progress = None
        self._task_id = None
        self._tqdm = None
        self._start_ts = time.monotonic()

        if _HAS_RICH_PROGRESS and sys.stdout.isatty():
            try:
                self._progress = Progress(
                    SpinnerColumn(
                        spinner_name=get_rich_spinner_name(),
                        style="cyan",
                        finished_text="[green]✔︎[/green]",
                    ),
                    TextColumn("[green]{task.description}"),
                    BarColumn(),
                    TextColumn("{task.percentage:>6.1f}%"),
                    DownloadColumn(),
                    TransferSpeedColumn(),
                    TimeElapsedColumn(),
                    TimeRemainingColumn(),
                    transient=True,
                )
                self._progress.start()
                self._task_id = self._progress.add_task(
                    self.desc,
                    total=self.total,
                )
                self._backend = "rich"
            except Exception:
                self._progress = None
                self._task_id = None

        if self._backend == "none" and _HAS_TQDM:
            self._tqdm = tqdm(
                total=self.total,
                desc=self.desc,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
                leave=False,
                bar_format="{desc}: {percentage:3.0f}%|{bar}| "
                           "[{elapsed}<{remaining}, {rate_fmt}{postfix}]",
            )
            self._backend = "tqdm"

    def update(self, n_bytes: int) -> None:
        step = int(max(0, n_bytes))
        if step == 0:
            return
        if self._backend == "rich" and self._progress is not None and self._task_id is not None:
            self._progress.update(self._task_id, advance=step)
        elif self._backend == "tqdm" and self._tqdm is not None:
            self._tqdm.update(step)

    def close(self, *, success: bool = True) -> None:
        elapsed_text = format_elapsed(time.monotonic() - self._start_ts)
        if self._backend == "rich" and self._progress is not None:
            if success and self._task_id is not None and self.total is not None:
                try:
                    self._progress.update(self._task_id, completed=int(self.total))
                except Exception:
                    pass
            self._progress.stop()
            self._progress = None
            self._task_id = None
        elif self._backend == "tqdm" and self._tqdm is not None:
            self._tqdm.close()
            self._tqdm = None
        if success:
            print_success(f"{self.desc} ...Finished [{elapsed_text}]")
        else:
            print_failure(f"{self.desc} ...Failed [{elapsed_text}]")


def downloader(url: str, dst: Path) -> None:
    dst = Path(dst)
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_suffix(dst.suffix + ".tmp")

    with urllib.request.urlopen(url) as r, open(tmp, "wb") as f:
        total = r.headers.get("Content-Length")
        total = int(total) if total is not None else None
        pbar = _DownloadProgress(total=total, desc="Downloading JanusX image")
        ok = False
        try:
            while True:
                chunk = r.read(1024 * 1024)
                if not chunk:
                    break
                f.write(chunk)
                pbar.update(len(chunk))
            ok = True
        finally:
            pbar.close(success=ok)

    os.replace(tmp, dst)


def _fastq_read_info(name: str):
    m = READ_RE.search(name)
    if not m:
        return None
    read = m.group(1).upper()
    stem = name[:m.start()]
    parts = stem.split("_")
    sample = parts[0]
    lane = parts[1] if len(parts) > 1 else None
    return sample, lane, read


def two_level_cluster(fastqlist):
    level1 = defaultdict(list)
    meta = {}

    for f in fastqlist:
        f = Path(f)
        if not f.name.lower().endswith(FASTQ_SUFFIXES):
            continue
        info = _fastq_read_info(f.name)
        if not info:
            continue
        sample, lane, read = info
        level1[sample].append(f)
        meta[f] = (sample, lane, read)

    level2 = {}
    for sample, files in level1.items():
        if len(files) == 4:
            sub = defaultdict(list)
            for f in files:
                s, lane, _ = meta[f]
                sub[(s, lane)].append(f)
            level2[sample] = dict(sub)

    return dict(level1), level2


def _pair_fastqs(files, sample_hint: str) -> tuple[Path, Path]:
    r1 = None
    r2 = None
    for f in sorted(files):
        info = _fastq_read_info(f.name)
        if not info:
            continue
        read = info[2]
        if read == "R1":
            if r1 is not None:
                raise ValueError(f"Duplicate R1 for sample {sample_hint}: {f}")
            r1 = f
        elif read == "R2":
            if r2 is not None:
                raise ValueError(f"Duplicate R2 for sample {sample_hint}: {f}")
            r2 = f

    if r1 is None or r2 is None:
        raise ValueError(f"Missing R1/R2 for sample {sample_hint}: {files}")
    return r1, r2


def build_samples(level1, level2):
    samples = {}
    for sample, files in level1.items():
        if len(files) == 2:
            r1, r2 = _pair_fastqs(files, sample)
            samples[sample] = [r1, r2]
            continue

        if len(files) == 4 and sample in level2:
            for (s, lane), lane_files in level2[sample].items():
                r1, r2 = _pair_fastqs(lane_files, f"{s}_{lane}")
                key = f"{s}_{lane}" if lane else s
                samples[key] = [r1, r2]
            continue

        raise ValueError(
            f"Sample {sample} has {len(files)} FASTQ files. "
            "Expected 2 (R1/R2) or 4 (two lanes) files."
        )
    return samples


def read_chroms_from_fai(reference: Path) -> List[str]:
    fai = Path(f"{reference}.fai")
    if not fai.exists():
        raise FileNotFoundError(f"FASTA index not found: {fai}")
    chroms: List[str] = []
    with open(fai, "r", encoding="utf-8") as f:
        for line in f:
            if not line.strip():
                continue
            chrom = line.split("\t", 1)[0].strip()
            if chrom:
                chroms.append(chrom)
    if not chroms:
        raise ValueError(f"Empty FASTA index: {fai}")
    return chroms


def main():
    t_start = time.time()
    use_spinner = bool(getattr(sys.stdout, "isatty", lambda: False)())

    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )

    required_group = parser.add_argument_group("Required Arguments")
    required_group.add_argument(
        "-r", "--reference", required=True, type=str,
        help="Reference genome FASTA file.",
    )
    required_group.add_argument(
        "-i", "--fastq-dir", required=True, type=str,
        help="Directory containing FASTQ files.",
    )
    required_group.add_argument(
        "-w", "--workdir", required=True, type=str,
        help="Working directory for outputs.",
    )
    optional_group = parser.add_argument_group("Optional Arguments")
    optional_group.add_argument(
        "-b", "--backend", required=False, type=str,default='nohup',
        choices=["nohup", "csub"],
        help=f"Execution backend for the pipeline. (default: nohup)",
    )
    optional_group.add_argument(
        "-m", "--maxtask", required=False, type=int,default=1,
        help="Max concurrent jobs when backend=nohup. (default: 1)",
    )

    args = parser.parse_args()

    workdir = Path(args.workdir).expanduser().resolve()
    workdir.mkdir(parents=True, exist_ok=True)
    (workdir / "log").mkdir(parents=True, exist_ok=True)
    (workdir / "tmp").mkdir(parents=True, exist_ok=True)
    log_path = workdir / "fastq2vcf.log"
    logger = setup_logging(str(log_path))

    reference = Path(args.reference).expanduser().resolve()
    if not ensure_file_exists(logger, str(reference), "Reference file"):
        raise SystemExit(1)

    fastq_dir = Path(args.fastq_dir).expanduser().resolve()
    if not ensure_dir_exists(logger, str(fastq_dir), "FASTQ directory"):
        raise SystemExit(1)

    emit_cli_configuration(
        logger,
        app_title="JanusX - Fastq2VCF",
        config_title="FASTQ2VCF CONFIG",
        host=socket.gethostname(),
        sections=[
            (
                "General",
                [
                    ("Reference", reference),
                    ("FASTQ dir", fastq_dir),
                    ("Workdir", workdir),
                    ("Backend", args.backend),
                    ("Max tasks (nohup)", args.maxtask),
                ],
            )
        ],
        line_max_chars=60,
    )

    binary_path = Path(janusx.pipeline.__file__).resolve().parent / "bin"
    binary_path.mkdir(parents=True, exist_ok=True)
    sif_path = binary_path / "janusxext.sif"
    if not sif_path.exists():
        user_input = input(
            "Download JanusX Singularity Image (~800MB)? [y/N or /path/to.sif]: "
        ).strip()
        if not user_input or user_input.lower() == "n":
            raise SystemExit(1)
        tmp_path = binary_path / "janusxext.tmp.sif"
        if user_input.lower() == "y":
            downloader(
                "https://github.com/FJingxian/JanusX/releases/download/v1.0.10/janusxext.sif",
                tmp_path,
            )
        else:
            src = Path(user_input).expanduser().resolve()
            if not src.exists():
                logger.error(f"SIF path not found: {src}")
                raise SystemExit(1)
            shutil.copy2(src, tmp_path)

        # subprocess.run([str(tmp_path), "gatk", "-version"], shell=True, check=True)
        os.replace(tmp_path, sif_path)
        logger.info(f'Saved in {sif_path}')

    singularity_prefix = f"singularity exec {sif_path} "
    logger.info(f"Singularity: {singularity_prefix}")

    fai_path = Path(f"{reference}.fai")
    ann_path = Path(f"{reference}.ann")
    if not fai_path.exists() or not ann_path.exists():
        cmd = indexREF(str(reference), singularity=singularity_prefix)
        logger.info(f"Indexing reference: {cmd}")
        index_job = wrap_cmd(cmd, "indexREF", 1, scheduler=args.backend)
        with CliStatus("Indexing reference...", enabled=use_spinner) as task:
            try:
                subprocess.run(index_job, shell=True, check=True, cwd=workdir)
            except Exception:
                task.fail("Reference indexing ...Failed")
                raise
            task.complete("Reference indexing started")
        with CliStatus("Waiting for FASTA index files...", enabled=use_spinner) as task:
            while (not fai_path.exists()) or (not ann_path.exists()):
                time.sleep(5)
            task.complete("Reference index files ready")

    chrom = read_chroms_from_fai(reference)
    logger.info(f"Chromosomes: {len(chrom)}")

    fastqlist = sorted(
        p for p in fastq_dir.rglob("*")
        if p.is_file() and p.name.lower().endswith(FASTQ_SUFFIXES)
    )
    if not fastqlist:
        logger.error(f"No FASTQ files found in {fastq_dir}")
        raise SystemExit(1)

    level1, level2 = two_level_cluster(fastqlist)
    samples = build_samples(level1, level2)
    logger.info(f"Samples: {len(samples)}")

    metadata = {
        "reference": str(reference),
        "samples": {k: [str(v[0]), str(v[1])] for k, v in samples.items()},
        "chrom": chrom,
    }

    os.chdir(workdir)
    fastq2vcf(
        metadata=metadata,
        workdir=workdir,
        backbend=args.backend,
        nohup_max_jobs=args.maxtask,
        singularity=singularity_prefix,
    )

    logger.info(f"Pipeline finished in {time.time() - t_start:.1f}s")


if __name__ == "__main__":
    main()
