"""Audit protein FASTA under an explicit 20-residue policy; never rewrite input."""
import argparse
import csv
import hashlib
import sys
from collections import Counter
from pathlib import Path

AA = set("ACDEFGHIKLMNPQRSTVWY")


def records(path):
    header, chunks = None, []
    with path.open(encoding="utf-8-sig") as handle:
        for line_no, raw in enumerate(handle, 1):
            line = raw.rstrip("\r\n")
            if not line.strip():
                continue
            if line.startswith(">"):
                if header is not None:
                    yield header, "".join(chunks)
                header, chunks = line[1:].strip(), []
                if not header:
                    raise ValueError(f"line {line_no}: empty FASTA header")
            else:
                if header is None:
                    raise ValueError(f"line {line_no}: sequence before first header")
                chunks.append(line)
    if header is not None:
        yield header, "".join(chunks)


def audit(path, max_residues):
    rows, first_sequence = [], {}
    for number, (header, raw) in enumerate(records(path), 1):
        ident = header.split()[0]
        sequence = raw.upper()
        problems = []
        if not sequence:
            problems.append("empty_sequence")
        if any(c.isspace() for c in raw):
            problems.append("whitespace_in_sequence")
        invalid = sorted(set(raw) - AA - set("acdefghiklmnpqrstvwy"))
        if invalid:
            problems.append("unsupported=" + repr("".join(invalid)))
        if len(sequence) > max_residues:
            problems.append("over_limit")
        # Hash and duplicate comparison only for canonical, nonempty sequences.
        canonical = bool(sequence) and not invalid
        digest = hashlib.sha256(sequence.encode("ascii")).hexdigest() if canonical else ""
        duplicate = ""
        if canonical:
            duplicate = first_sequence.get(sequence, "")
            first_sequence.setdefault(sequence, number)
        rows.append(dict(record=number, id=ident, length=len(sequence),
                         lowercase=raw != sequence, sha256=digest,
                         duplicate_of_record=duplicate, issues=";".join(problems)))
    if not rows:
        raise ValueError("no FASTA records")
    counts = Counter(row["id"] for row in rows)
    for row in rows:
        if counts[row["id"]] > 1:
            row["issues"] = ";".join(filter(None, [row["issues"], "duplicate_id"]))
    return rows


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("fasta", type=Path)
    parser.add_argument("--max-residues", type=int, required=True)
    parser.add_argument("--out", type=Path, default=Path("audit.csv"))
    args = parser.parse_args()
    if args.max_residues < 1:
        parser.error("--max-residues must be positive")
    try:
        rows = audit(args.fasta, args.max_residues)
        with args.out.open("x", newline="", encoding="utf-8") as handle:
            writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
            writer.writeheader()
            writer.writerows(rows)
    except (OSError, ValueError) as error:
        parser.exit(2, f"error: {error}\n")
    flagged = sum(bool(row["issues"]) for row in rows)
    repeats = sum(row["duplicate_of_record"] != "" for row in rows)
    print(f"records={len(rows)} flagged={flagged} exact_repeats={repeats}")
    print(f"report={args.out}")
    return 1 if flagged or repeats else 0


if __name__ == "__main__":
    sys.exit(main())