"""BABAPPA command-line interface."""

from __future__ import annotations

import argparse
import json
import platform
import shutil
import sys
from pathlib import Path

import pandas as pd

from babappa.benchmark.runners import run_all_benchmarks
from babappa.config import BABAPPAConfig, default_config, load_config
from babappa.io.fasta import read_alignment_map, read_codon_alignment
from babappa.io.newick import read_newick
from babappa.io.project import init_project
from babappa.model.bundle import read_bundle
from babappa.qc.validators import validate_alignment_tree, write_qc_results
from babappa.report.html import build_report
from babappa.utils.fs import ensure_dir
from babappa.utils.logging import configure_logging
from babappa.utils.manifest import build_manifest, write_manifest
from babappa.utils.package_data import example_data_dir
from babappa.version import __version__
from babappa.workflow import (
    adjust_scores_table,
    fit_null_bundle,
    run_sensitivity_analysis,
    score_genes,
)


def _parse_csv_floats(text: str | None) -> list[float]:
    if not text:
        return []
    return [float(x.strip()) for x in text.split(",") if x.strip()]


def _parse_csv_ints(text: str | None) -> list[int]:
    if not text:
        return []
    return [int(x.strip()) for x in text.split(",") if x.strip()]


def _parse_csv_strings(text: str | None) -> list[str]:
    if not text:
        return []
    return [x.strip() for x in text.split(",") if x.strip()]


def _load_cfg(path: str | None) -> BABAPPAConfig:
    return load_config(path) if path else default_config()


def cmd_init(args: argparse.Namespace) -> int:
    out = init_project(args.path)
    print(f"Initialized BABAPPA project at: {out}")
    return 0


def cmd_validate(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "validate.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    alignments = read_alignment_map(single=args.alignment, directory=args.alignment_dir)
    tree = read_newick(args.tree)

    results = [validate_alignment_tree(aln, tree, cfg.qc) for aln in alignments.values()]
    tsv, js = write_qc_results(results, outdir)
    manifest = build_manifest(
        command="validate",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={"qc_summary_tsv": str(tsv), "qc_summary_json": str(js)},
    )
    write_manifest(outdir / "manifest.json", manifest)

    n_err = sum(0 if r.ok else 1 for r in results)
    if n_err:
        logger.error("Validation completed with %d gene(s) failing QC.", n_err)
        for r in results:
            if r.errors:
                logger.error("%s errors: %s", r.gene, " | ".join(r.errors))
        return 2

    logger.info("Validation completed successfully for %d gene(s)", len(results))
    return 0


def cmd_fit_null(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "fit_null.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    alignment = read_codon_alignment(args.alignment)
    tree = read_newick(args.tree)
    qc = validate_alignment_tree(alignment, tree, cfg.qc)
    if qc.errors:
        logger.error("Input QC failed: %s", " | ".join(qc.errors))
        return 2

    bundle, bundle_path, diag_df = fit_null_bundle(
        alignment=alignment,
        tree=tree,
        cfg=cfg,
        outdir=outdir,
        seed=args.seed,
        energy_family=args.energy_family,
        null_replicates_train=args.null_replicates_train,
    )

    manifest = build_manifest(
        command="fit-null",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={
            "bundle_path": str(bundle_path),
            "n_training_replicates": int(bundle.training["n_null_replicates_train"]),
            "n_training_sites": int(bundle.training["n_training_sites"]),
        },
    )
    write_manifest(outdir / "manifest.json", manifest)
    logger.info("Frozen null bundle written to %s", bundle_path)
    logger.info("Training diagnostics: %s", outdir / "null_training_diagnostics.tsv")
    logger.info("Median training gene statistic: %.4f", float(diag_df["statistic"].median()))
    return 0


def cmd_score(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "score.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    alignments = read_alignment_map(single=args.alignment, directory=args.alignment_dir)
    tree = read_newick(args.tree)
    bundle = read_bundle(args.bundle)

    # QC pass before scoring to provide readable errors.
    failures = []
    for gene, aln in alignments.items():
        qc = validate_alignment_tree(aln, tree, cfg.qc)
        if qc.errors:
            failures.append((gene, qc.errors))
    if failures:
        for gene, errs in failures:
            logger.error("QC failed for %s: %s", gene, " | ".join(errs))
        return 2

    scores = score_genes(
        alignments=alignments,
        tree=tree,
        bundle=bundle,
        cfg=cfg,
        outdir=outdir,
        seed=args.seed,
        jobs=args.jobs,
        resume=args.resume,
        n_replicates_score=args.null_replicates_score,
        two_sided=args.two_sided,
    )

    sensitivity_path = None
    if args.sensitivity:
        run_sensitivity_analysis(
            alignments=alignments,
            tree=tree,
            base_bundle=bundle,
            cfg=cfg,
            outdir=outdir / "sensitivity",
            seed=args.seed + 701,
            bandwidths=_parse_csv_floats(args.sensitivity_bandwidths),
            kappas=_parse_csv_floats(args.sensitivity_kappas),
            replicate_counts=_parse_csv_ints(args.sensitivity_replicates),
            energy_families=_parse_csv_strings(args.sensitivity_energy_families),
            codon_frequency_modes=_parse_csv_strings(args.sensitivity_codon_frequencies),
        )
        sensitivity_path = outdir / "sensitivity" / "sensitivity_summary.tsv"
        logger.info("Sensitivity summary saved to %s", sensitivity_path)

    manifest = build_manifest(
        command="score",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={
            "scores_tsv": str(outdir / "scores.tsv"),
            "n_genes": int(scores.shape[0]),
            "sensitivity_summary": str(sensitivity_path) if sensitivity_path else None,
        },
    )
    write_manifest(outdir / "manifest.json", manifest)

    logger.info("Scored %d gene(s). Output: %s", scores.shape[0], outdir / "scores.tsv")
    return 0


def cmd_adjust(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "adjust.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    scores = pd.read_csv(args.scores, sep="\t")
    if "p_value" not in scores.columns:
        raise ValueError("Input score table must contain a 'p_value' column")
    adjusted = adjust_scores_table(scores, method=args.method)

    out_tsv = outdir / "scores_adjusted.tsv"
    adjusted.to_csv(out_tsv, sep="\t", index=False)

    manifest = build_manifest(
        command="adjust",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={"adjusted_scores_tsv": str(out_tsv), "method": args.method},
    )
    write_manifest(outdir / "manifest.json", manifest)
    logger.info("Adjusted scores written to %s", out_tsv)
    return 0


def cmd_report(args: argparse.Namespace) -> int:
    from babappa.plots.figures import (
        plot_null_distribution,
        plot_sensitivity,
        plot_site_contributions,
    )

    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "report.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    scores = pd.read_csv(args.scores, sep="\t")
    if scores.empty:
        raise ValueError("Scores table is empty")

    source_score_dir = Path(args.score_dir) if args.score_dir else Path(args.scores).resolve().parent
    first_gene = str(scores.iloc[0]["gene"])
    null_json = source_score_dir / "genes" / first_gene / "null_distribution.json"
    site_table = source_score_dir / "genes" / first_gene / "site_contributions.tsv"

    plots: dict[str, str] = {"null_distribution": "", "site_contribution": "", "sensitivity": ""}
    warnings: list[str] = []

    if null_json.exists():
        payload = json.loads(null_json.read_text(encoding="utf-8"))
        p = plot_null_distribution(
            observed=float(payload["observed_statistic"]),
            null_stats=pd.Series(payload["null_statistics"]).to_numpy(),
            outpath=outdir / "null_distribution.png",
        )
        plots["null_distribution"] = p.name
        warnings.extend(payload.get("inference", {}).get("warning_flags", []))
        warnings.extend(payload.get("qc", {}).get("flags", []))

    if site_table.exists():
        site_df = pd.read_csv(site_table, sep="\t")
        p = plot_site_contributions(site_df, outpath=outdir / "site_contributions.png")
        plots["site_contribution"] = p.name

    sens = source_score_dir / "sensitivity" / "sensitivity_summary.tsv"
    if sens.exists():
        sens_df = pd.read_csv(sens, sep="\t")
        p = plot_sensitivity(sens_df, outpath=outdir / "sensitivity.png")
        plots["sensitivity"] = p.name

    manifest_path = str((source_score_dir / "manifest.json").resolve()) if (source_score_dir / "manifest.json").exists() else "N/A"
    html_path = build_report(
        scores=scores,
        outpath=outdir / "report.html",
        manifest_path=manifest_path,
        null_model_desc=args.null_model_description,
        smoothing_bandwidth=cfg.representation.smoothing_bandwidth,
        warnings=sorted(set(warnings)),
        plots=plots,
    )

    manifest = build_manifest(
        command="report",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={"report_html": str(html_path)},
    )
    write_manifest(outdir / "manifest.json", manifest)
    logger.info("Report written to %s", html_path)
    return 0


def cmd_benchmark(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "benchmark.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    results = run_all_benchmarks(cfg=cfg, outdir=outdir, seed=args.seed)

    manifest = build_manifest(
        command="benchmark",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={
            "simulation_summary": str(outdir / "simulation" / "simulation_summary.tsv"),
            "real_data_summary": str(outdir / "real_data" / "real_data_summary.tsv"),
        },
    )
    write_manifest(outdir / "manifest.json", manifest)
    logger.info("Benchmark complete. Tracks: %s", ", ".join(results.keys()))
    return 0


def cmd_demo(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "demo.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    data = example_data_dir() / "small_demo"
    aln_dir = data / "alignments"
    tree = data / "tree.nwk"

    # Validate
    val_out = outdir / "validate"
    validate_args = argparse.Namespace(
        alignment=None,
        alignment_dir=str(aln_dir),
        tree=str(tree),
        outdir=str(val_out),
        config=args.config,
        seed=args.seed,
        verbose=args.verbose,
    )
    code = cmd_validate(validate_args)
    if code != 0:
        return code

    alignments = read_alignment_map(directory=str(aln_dir))
    ref_gene = sorted(alignments.keys())[0]

    fit_out = outdir / "fit_null"
    bundle, bundle_path, _ = fit_null_bundle(
        alignment=alignments[ref_gene],
        tree=read_newick(tree),
        cfg=cfg,
        outdir=fit_out,
        seed=args.seed,
        energy_family="gaussian",
    )

    score_out = outdir / "score"
    scores = score_genes(
        alignments=alignments,
        tree=read_newick(tree),
        bundle=bundle,
        cfg=cfg,
        outdir=score_out,
        seed=args.seed + 11,
        jobs=max(1, args.jobs),
        resume=False,
    )

    adj_out = outdir / "adjust"
    adjusted = adjust_scores_table(scores, method="bh")
    adj_out.mkdir(parents=True, exist_ok=True)
    adjusted.to_csv(adj_out / "scores_adjusted.tsv", sep="\t", index=False)

    rep_out = outdir / "report"
    report_args = argparse.Namespace(
        scores=str(adj_out / "scores_adjusted.tsv"),
        outdir=str(rep_out),
        score_dir=str(score_out),
        config=args.config,
        seed=args.seed,
        verbose=args.verbose,
        null_model_description="MG94-like neutral codon model on supplied tree",
    )
    cmd_report(report_args)

    manifest = build_manifest(
        command="demo",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={"bundle": str(bundle_path), "report": str(rep_out / "report.html")},
    )
    write_manifest(outdir / "manifest.json", manifest)
    logger.info("Demo complete. Open %s", rep_out / "report.html")
    return 0


def cmd_doctor(args: argparse.Namespace) -> int:
    outdir = ensure_dir(args.outdir)
    logger = configure_logging(outdir / "doctor.log", verbose=args.verbose)
    cfg = _load_cfg(args.config)

    checks = {
        "python": sys.version,
        "platform": platform.platform(),
        "babappa_version": __version__,
        "numpy": __import__("numpy").__version__,
        "pandas": __import__("pandas").__version__,
        "scipy": __import__("scipy").__version__,
        "matplotlib": __import__("matplotlib").__version__,
        "jinja2": __import__("jinja2").__version__,
        "yaml": __import__("yaml").__version__,
        "codeml_available": shutil.which("codeml") is not None,
        "hyphy_available": shutil.which("hyphy") is not None,
    }

    (outdir / "doctor.json").write_text(json.dumps(checks, indent=2), encoding="utf-8")
    manifest = build_manifest(
        command="doctor",
        config=cfg,
        seed=args.seed,
        args=vars(args),
        extras={"doctor_json": str(outdir / "doctor.json")},
    )
    write_manifest(outdir / "manifest.json", manifest)

    logger.info("BABAPPA doctor completed. See %s", outdir / "doctor.json")
    for k, v in checks.items():
        logger.info("%s: %s", k, v)
    return 0


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="babappa",
        description=(
            "BABAPPA: gene-level detection of excess evolutionary dispersion relative to an explicit neutral model"
        ),
    )
    p.add_argument("--version", action="version", version=f"babappa {__version__}")

    sub = p.add_subparsers(dest="command", required=True)

    init = sub.add_parser("init", help="Create a new BABAPPA project scaffold")
    init.add_argument("path", help="Project directory to create")
    init.set_defaults(func=cmd_init)

    def add_common_io(ap: argparse.ArgumentParser) -> None:
        group = ap.add_mutually_exclusive_group(required=True)
        group.add_argument("--alignment", help="Single codon alignment FASTA")
        group.add_argument("--alignment-dir", help="Directory with codon alignment FASTA files")
        ap.add_argument("--tree", required=True, help="Newick tree path")
        ap.add_argument("--config", help="YAML/TOML config file")
        ap.add_argument("--outdir", required=True, help="Output directory")
        ap.add_argument("--seed", type=int, default=1, help="Random seed")
        ap.add_argument("--verbose", action="store_true", help="Verbose logging")

    validate = sub.add_parser("validate", help="Validate alignments and tree consistency")
    add_common_io(validate)
    validate.set_defaults(func=cmd_validate)

    fit = sub.add_parser("fit-null", help="Simulate neutral data, fit and freeze null energy model")
    fit.add_argument("--alignment", required=True, help="Reference codon alignment FASTA")
    fit.add_argument("--tree", required=True, help="Newick tree path")
    fit.add_argument("--config", help="YAML/TOML config file")
    fit.add_argument("--outdir", required=True, help="Output directory")
    fit.add_argument("--seed", type=int, default=1, help="Random seed")
    fit.add_argument("--null-replicates-train", type=int, help="Number of null simulations for fitting")
    fit.add_argument(
        "--energy-family",
        default="gaussian",
        choices=["gaussian", "diagonal_gaussian"],
        help="Energy model family",
    )
    fit.add_argument("--verbose", action="store_true", help="Verbose logging")
    fit.set_defaults(func=cmd_fit_null)

    score = sub.add_parser("score", help="Score genes with frozen null model and exact Monte Carlo p-values")
    add_common_io(score)
    score.add_argument("--bundle", required=True, help="Frozen null bundle JSON from fit-null")
    score.add_argument("--jobs", type=int, default=1, help="Parallel jobs for per-gene scoring")
    score.add_argument("--resume", action="store_true", help="Resume from existing per-gene outputs")
    score.add_argument("--null-replicates-score", type=int, help="Null replicates for Monte Carlo p-values")
    score.add_argument("--two-sided", action="store_true", help="Use pooled symmetric two-sided p-value")
    score.add_argument("--sensitivity", action="store_true", help="Run sensitivity analysis")
    score.add_argument("--sensitivity-bandwidths", help="Comma-separated smoothing bandwidths")
    score.add_argument("--sensitivity-kappas", help="Comma-separated kappa values")
    score.add_argument("--sensitivity-replicates", help="Comma-separated replicate counts")
    score.add_argument("--sensitivity-energy-families", help="Comma-separated energy families")
    score.add_argument("--sensitivity-codon-frequencies", help="Comma-separated codon frequency modes")
    score.set_defaults(func=cmd_score)

    adjust = sub.add_parser("adjust", help="Apply multiple-testing correction")
    adjust.add_argument("--scores", required=True, help="Input score TSV")
    adjust.add_argument("--outdir", required=True, help="Output directory")
    adjust.add_argument("--method", default="bh", choices=["bh", "by", "bonferroni"], help="Adjustment method")
    adjust.add_argument("--config", help="YAML/TOML config file")
    adjust.add_argument("--seed", type=int, default=1, help="Random seed")
    adjust.add_argument("--verbose", action="store_true", help="Verbose logging")
    adjust.set_defaults(func=cmd_adjust)

    report = sub.add_parser("report", help="Build HTML report with figures and interpretation")
    report.add_argument("--scores", required=True, help="Score/adjusted score TSV")
    report.add_argument("--outdir", required=True, help="Output directory")
    report.add_argument("--score-dir", help="Directory containing score artifacts (genes/*)")
    report.add_argument("--config", help="YAML/TOML config file")
    report.add_argument("--seed", type=int, default=1, help="Random seed")
    report.add_argument("--verbose", action="store_true", help="Verbose logging")
    report.add_argument(
        "--null-model-description",
        default="Configured neutral codon model on supplied tree",
        help="Text description for report",
    )
    report.set_defaults(func=cmd_report)

    bench = sub.add_parser("benchmark", help="Run simulation and real-data benchmark suites")
    bench.add_argument("--outdir", required=True, help="Output directory")
    bench.add_argument("--config", help="YAML/TOML config file")
    bench.add_argument("--seed", type=int, default=1, help="Random seed")
    bench.add_argument("--verbose", action="store_true", help="Verbose logging")
    bench.set_defaults(func=cmd_benchmark)

    demo = sub.add_parser("demo", help="Run complete minimal example from packaged data")
    demo.add_argument("--outdir", required=True, help="Output directory")
    demo.add_argument("--config", help="YAML/TOML config file")
    demo.add_argument("--seed", type=int, default=7, help="Random seed")
    demo.add_argument("--jobs", type=int, default=1, help="Parallel jobs")
    demo.add_argument("--verbose", action="store_true", help="Verbose logging")
    demo.set_defaults(func=cmd_demo)

    doctor = sub.add_parser("doctor", help="Check environment, dependencies, and optional tools")
    doctor.add_argument("--outdir", required=True, help="Output directory")
    doctor.add_argument("--config", help="YAML/TOML config file")
    doctor.add_argument("--seed", type=int, default=1, help="Random seed")
    doctor.add_argument("--verbose", action="store_true", help="Verbose logging")
    doctor.set_defaults(func=cmd_doctor)

    return p


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return int(args.func(args))
    except Exception as exc:  # pragma: no cover
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1


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