#!/usr/bin/env -S uv run
# /// script
# dependencies = []
# ///
#

# uv run prepare_corpus.py prepare ../data ./corpus.txt --recursive


import argparse
import re
import unicodedata
from pathlib import Path


def clean_text(text: str) -> str:
    # Unicode normalization
    text = unicodedata.normalize("NFC", text)

    # Remove control characters except newline/tab
    text = "".join(
        ch for ch in text if unicodedata.category(ch)[0] != "C" or ch in "\n\t"
    )

    # Normalize line endings
    text = text.replace("\r\n", "\n").replace("\r", "\n")

    # Remove zero-width characters (common OCR artifacts)
    text = re.sub(r"[\u200B-\u200D\uFEFF]", "", text)

    # Normalize spaces
    text = re.sub(r"[ \t]+", " ", text)

    # Remove common OCR symbols
    text = re.sub(r"[■□▪▫◆◇★☆◦•●]+", " ", text)
    text = re.sub(r"[_=~`]+", " ", text)

    # Remove separator lines
    text = re.sub(r"(?m)^[-=*_]{3,}\s*$", "", text)

    # Remove standalone page numbers
    text = re.sub(r"(?m)^\s*\d+\s*$", "", text)

    # Remove lines containing only Roman numerals (optional)
    text = re.sub(r"(?m)^\s*[IVXLCDM]+\s*$", "", text)

    # Trim spaces around newlines
    text = re.sub(r" *\n *", "\n", text)

    # Collapse multiple blank lines
    text = re.sub(r"\n{3,}", "\n\n", text)

    return text.strip()


def collect_text_files(folder: Path, recursive: bool):
    pattern = "**/*.txt" if recursive else "*.txt"
    return sorted(folder.glob(pattern))


def build_corpus(files, output_file: Path):
    print(f"output_file {output_file}")

    with output_file.open("w", encoding="utf-8") as out:
        for file in files:
            print(file)

            try:
                text = file.read_text(
                    encoding="utf-8",
                    errors="ignore",
                )

                text = clean_text(text)

                out.write(text)
                out.write("\n")

            except Exception as e:
                print(f"Skipping {file}: {e}")

    print("DONE build_corpus")


def prepare_command(args):
    files = collect_text_files(Path(args.folder), args.recursive)
    if not files:
        raise SystemExit("No .txt files found")
    build_corpus(files, Path(args.output))
    print(f"Corpus written to {args.output}")


def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="command", required=True)

    p = sub.add_parser("prepare")
    p.add_argument("folder", default="../data/")
    p.add_argument("output", default="./corpus.txt")
    p.add_argument("--recursive", action="store_true")

    args = parser.parse_args()

    if args.command == "prepare":
        prepare_command(args)


if __name__ == "__main__":
    main()