"""Input QC validation for codon alignments and trees."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd

from babappa.config import QCConfig
from babappa.io.fasta import CodonAlignment
from babappa.io.newick import PhyloTree, tree_summary
from babappa.utils.fs import write_json

STOP_CODONS = {"TAA", "TAG", "TGA"}
VALID_NT = {"A", "C", "G", "T", "N", "-"}


@dataclass
class QCResult:
    gene: str
    metrics: dict[str, float | int]
    warnings: list[str]
    errors: list[str]
    flags: list[str]

    @property
    def ok(self) -> bool:
        return not self.errors

    def to_dict(self) -> dict[str, Any]:
        return {
            "gene": self.gene,
            "metrics": self.metrics,
            "warnings": self.warnings,
            "errors": self.errors,
            "flags": self.flags,
            "ok": self.ok,
        }


def _codons(seq: str) -> list[str]:
    return [seq[i : i + 3] for i in range(0, len(seq), 3)]


def _is_ambiguous(codon: str) -> bool:
    return any(base not in {"A", "C", "G", "T", "-"} for base in codon)


def validate_alignment_tree(alignment: CodonAlignment, tree: PhyloTree, qc_cfg: QCConfig) -> QCResult:
    warnings: list[str] = []
    errors: list[str] = []
    flags: list[str] = []

    seqs = alignment.sequences
    taxa_aln = list(seqs.keys())
    taxa_tree = tree.leaf_names()

    if len(taxa_aln) != len(set(taxa_aln)):
        errors.append("Duplicate taxa in alignment")

    missing_from_tree = sorted(set(taxa_aln) - set(taxa_tree))
    missing_from_alignment = sorted(set(taxa_tree) - set(taxa_aln))
    if missing_from_tree:
        errors.append(f"Alignment taxa missing in tree: {', '.join(missing_from_tree)}")
    if missing_from_alignment:
        errors.append(f"Tree taxa missing in alignment: {', '.join(missing_from_alignment)}")

    lengths = [len(seq) for seq in seqs.values()]
    if len(set(lengths)) != 1:
        errors.append("Alignment sequences have inconsistent lengths")
        nt_len = min(lengths) if lengths else 0
    else:
        nt_len = lengths[0] if lengths else 0

    if nt_len == 0:
        errors.append("Empty sequences")

    non_triplet = sum(1 for n in lengths if n % 3 != 0)
    if non_triplet > 0:
        errors.append(f"{non_triplet} sequences have non-triplet length")

    invalid_chars = 0
    gap_chars = 0
    ambiguous_chars = 0
    internal_stop_sites = 0
    low_info_taxa = 0

    codon_sites = nt_len // 3 if nt_len else 0
    for _taxon, seq in seqs.items():
        for ch in seq:
            if ch.upper() not in VALID_NT:
                invalid_chars += 1
            if ch == "-":
                gap_chars += 1
            if ch.upper() not in {"A", "C", "G", "T", "-"}:
                ambiguous_chars += 1

        if codon_sites > 0:
            cds = _codons(seq.upper())
            informative = sum(1 for c in cds if c != "---" and not _is_ambiguous(c))
            if informative < max(3, int(0.1 * codon_sites)):
                low_info_taxa += 1
            for c in cds[:-1]:
                if c in STOP_CODONS:
                    internal_stop_sites += 1

    if invalid_chars > 0:
        errors.append(f"Invalid nucleotide characters found: {invalid_chars}")
    if internal_stop_sites > 0:
        warnings.append(f"Internal stop codons detected: {internal_stop_sites}")
        flags.append("internal_stop_codons")
    if low_info_taxa > 0:
        warnings.append(f"Low-information sequences: {low_info_taxa}")
        flags.append("low_taxon_information")

    total_chars = max(1, nt_len * max(1, len(seqs)))
    gap_fraction = gap_chars / total_chars
    ambiguity_fraction = ambiguous_chars / total_chars

    if gap_fraction > qc_cfg.max_gap_fraction:
        warnings.append(
            f"Gap burden {gap_fraction:.3f} exceeds threshold {qc_cfg.max_gap_fraction:.3f}"
        )
        flags.append("high_gap_burden")
    if ambiguity_fraction > qc_cfg.max_ambiguity_fraction:
        warnings.append(
            f"Ambiguity burden {ambiguity_fraction:.3f} exceeds threshold {qc_cfg.max_ambiguity_fraction:.3f}"
        )
        flags.append("high_ambiguity_burden")

    effective_sites = 0
    if codon_sites > 0:
        ordered_taxa = sorted(taxa_aln)
        for idx in range(codon_sites):
            site = [seqs[t][idx * 3 : idx * 3 + 3] for t in ordered_taxa]
            informative = sum(1 for c in site if c != "---" and not _is_ambiguous(c))
            if informative >= 2:
                effective_sites += 1

    if effective_sites < qc_cfg.min_effective_codon_sites:
        warnings.append(
            f"Effective informative codon sites ({effective_sites}) below threshold "
            f"{qc_cfg.min_effective_codon_sites}"
        )
        flags.append("low_effective_length")

    tree_stats = tree_summary(tree)
    min_nonzero = float(tree_stats["min_nonzero_branch_length"])
    max_len = float(tree_stats["max_branch_length"])
    branch_ratio = float("inf") if min_nonzero == 0 else max_len / min_nonzero
    if branch_ratio > qc_cfg.max_branch_imbalance_ratio:
        warnings.append(
            f"Branch-length imbalance ratio {branch_ratio:.2f} exceeds "
            f"threshold {qc_cfg.max_branch_imbalance_ratio:.2f}"
        )
        flags.append("branch_length_imbalance")

    metrics: dict[str, float | int] = {
        "n_taxa_alignment": len(taxa_aln),
        "n_taxa_tree": len(taxa_tree),
        "nt_length": nt_len,
        "codon_length": codon_sites,
        "gap_fraction": gap_fraction,
        "ambiguity_fraction": ambiguity_fraction,
        "effective_informative_sites": effective_sites,
        "internal_stop_sites": internal_stop_sites,
        "low_information_taxa": low_info_taxa,
        "branch_length_ratio": branch_ratio,
    }
    return QCResult(gene=alignment.gene, metrics=metrics, warnings=warnings, errors=errors, flags=flags)


def write_qc_results(results: list[QCResult], outdir: str | Path) -> tuple[Path, Path]:
    out_path = Path(outdir)
    out_path.mkdir(parents=True, exist_ok=True)

    table = pd.DataFrame([
        {
            "gene": r.gene,
            "ok": r.ok,
            "warnings": " | ".join(r.warnings),
            "errors": " | ".join(r.errors),
            "flags": "|".join(sorted(set(r.flags))),
            **r.metrics,
        }
        for r in results
    ])
    tsv = out_path / "qc_summary.tsv"
    table.to_csv(tsv, sep="\t", index=False)

    payload = {"results": [r.to_dict() for r in results]}
    json_path = out_path / "qc_summary.json"
    write_json(json_path, payload)
    return tsv, json_path
