"""FASTA parsing utilities for codon alignments."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path


@dataclass
class CodonAlignment:
    gene: str
    sequences: dict[str, str]

    @property
    def taxa(self) -> list[str]:
        return list(self.sequences.keys())

    @property
    def nt_length(self) -> int:
        lengths = {len(v) for v in self.sequences.values()}
        if len(lengths) != 1:
            raise ValueError(f"Inconsistent sequence lengths for gene {self.gene}: {sorted(lengths)}")
        return next(iter(lengths))

    @property
    def codon_length(self) -> int:
        return self.nt_length // 3


def read_fasta(path: str | Path) -> dict[str, str]:
    fasta_path = Path(path)
    if not fasta_path.exists():
        raise FileNotFoundError(f"FASTA not found: {fasta_path}")

    seqs: dict[str, str] = {}
    current: str | None = None
    chunks: list[str] = []

    for raw in fasta_path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line:
            continue
        if line.startswith(">"):  # header
            if current is not None:
                seqs[current] = "".join(chunks).upper().replace(" ", "")
            current = line[1:].strip().split()[0]
            if not current:
                raise ValueError(f"Empty FASTA header in {fasta_path}")
            if current in seqs:
                raise ValueError(f"Duplicate FASTA taxon '{current}' in {fasta_path}")
            chunks = []
        else:
            if current is None:
                raise ValueError(f"Sequence line before FASTA header in {fasta_path}")
            chunks.append(line)

    if current is not None:
        seqs[current] = "".join(chunks).upper().replace(" ", "")

    if not seqs:
        raise ValueError(f"No FASTA records found in {fasta_path}")
    return seqs


def read_codon_alignment(path: str | Path, gene_name: str | None = None) -> CodonAlignment:
    seqs = read_fasta(path)
    gene = gene_name or Path(path).stem
    return CodonAlignment(gene=gene, sequences=seqs)


def read_alignment_map(single: str | None = None, directory: str | None = None) -> dict[str, CodonAlignment]:
    if bool(single) == bool(directory):
        raise ValueError("Provide exactly one of --alignment or --alignment-dir")

    if single:
        aln = read_codon_alignment(single)
        return {aln.gene: aln}

    dir_path = Path(directory or "")
    if not dir_path.exists() or not dir_path.is_dir():
        raise FileNotFoundError(f"Alignment directory not found: {dir_path}")

    out: dict[str, CodonAlignment] = {}
    files = sorted([*dir_path.glob("*.fa"), *dir_path.glob("*.fasta"), *dir_path.glob("*.fas")])
    if not files:
        raise ValueError(f"No FASTA files found in {dir_path}")

    for file_path in files:
        aln = read_codon_alignment(file_path)
        if aln.gene in out:
            raise ValueError(f"Duplicate gene name from FASTA stem: {aln.gene}")
        out[aln.gene] = aln
    return out


def write_alignment(alignment: CodonAlignment, path: str | Path) -> None:
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    lines: list[str] = []
    for taxon, seq in alignment.sequences.items():
        lines.append(f">{taxon}")
        lines.append(seq)
    out.write_text("\n".join(lines) + "\n", encoding="utf-8")
