#!/usr/bin/env python3
"""
words-drill -- generate a personalized typing drill HTML page from weak bigrams/trigrams.

Usage:
  words-drill [--name NAME] [--corpus PATH] [--count N] [--out PATH] BIGRAM [BIGRAM ...]
  words-drill --stdin [other flags]   # read bigrams one per line from stdin

Bigrams are 2- or 3-character sequences (e.g. th, qu, tio). For symbol
bigrams (=>, ;;) pass --corpus pointing at a code wordlist.
"""

import argparse
import os
import random
import re
import sys
from datetime import date
from pathlib import Path

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

WORDS_ROOT = Path("/home/krisyotam/dev/words")
DEFAULT_CORPUS = WORDS_ROOT / "sources/google-10000-english/google-10000-english-usa-no-swears.txt"
DRILLS_DIR = WORDS_ROOT / "drills"

# ---------------------------------------------------------------------------
# HTML helpers (match build.py style exactly)
# ---------------------------------------------------------------------------

HEAD_TEMPLATE = """\
<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <meta name="author" content="Kris Yotam">
  <link rel="stylesheet" href="/words/style.css">
  <title>{title} -- Kris Yotam's Word Lists</title>
</head>
<body>
"""

FOOT = """\
</body>
</html>
"""


def render_drill_page(name: str, words: list, bigrams: list) -> str:
    bigram_list = ", ".join(bigrams)
    word_count = len(words)
    body = "\n".join(words)
    return (
        HEAD_TEMPLATE.format(title=name)
        + '<div class="text"><div class="left">\n'
        + '<a class="back-link" href="/words/drills/">&larr; back to drills</a>\n'
        + f"<h1>{name}</h1>\n"
        + f'<p class="wordlist-meta">{word_count:,} words &mdash; bigrams: {bigram_list}</p>\n'
        + f'<pre class="wordlist">{body}</pre>\n'
        + "</div></div>\n"
        + FOOT
    )


def render_drills_index(drill_files: list) -> str:
    """Render drills/index.html from a list of (stem, mtime) pairs, sorted by mtime desc."""
    listing_items = ""
    for stem, _mtime in drill_files:
        title = stem.replace("-", " ").title()
        listing_items += (
            '<div class="listing"><div class="left"><dl>'
            f'<dt><a href="{stem}.html">{title}</a></dt>'
            "</dl></div></div>\n"
        )

    return (
        HEAD_TEMPLATE.format(title="Drills")
        + '<div class="text"><div class="left">\n'
        + '<a class="back-link" href="/words/">&larr; back to word lists</a>\n'
        + "<h1>Drills</h1>\n"
        + "</div></div>\n"
        + listing_items
        + FOOT
    )


def render_drills_placeholder() -> str:
    return (
        HEAD_TEMPLATE.format(title="Drills")
        + '<div class="text"><div class="left">\n'
        + '<a class="back-link" href="/words/">&larr; back to word lists</a>\n'
        + "<h1>Drills</h1>\n"
        + '<p class="coming-soon">Coming soon -- LLM-generated symbol drills.</p>\n'
        + "</div></div>\n"
        + FOOT
    )


# ---------------------------------------------------------------------------
# Core logic
# ---------------------------------------------------------------------------

def load_corpus(path: Path) -> list:
    if not path.exists():
        print(f"error: corpus file not found: {path}", file=sys.stderr)
        sys.exit(1)
    lines = path.read_text(encoding="utf-8").splitlines()
    return [l.strip() for l in lines if l.strip()]


def find_words_for_bigram(corpus: list, bigram: str) -> list:
    """Return all corpus words that contain the given bigram."""
    return [w for w in corpus if bigram in w]


def build_drill_wordlist(corpus: list, bigrams: list, count: int, seed: int) -> tuple:
    """
    Distribute `count` words across bigrams proportionally.
    Returns (final_word_list, counts_dict).
    counts_dict maps each bigram to how many words were found/used.
    """
    per_bigram = count // len(bigrams) if bigrams else 0
    remainder = count - per_bigram * len(bigrams)

    rng = random.Random(seed)

    all_words = []
    counts = {}

    for i, bigram in enumerate(bigrams):
        matches = find_words_for_bigram(corpus, bigram)
        quota = per_bigram + (1 if i < remainder else 0)
        if not matches:
            counts[bigram] = 0
            continue
        rng.shuffle(matches)
        selected = matches[:quota]
        counts[bigram] = len(selected)
        all_words.extend(selected)

    # Shuffle final combined list with same rng for reproducibility
    rng.shuffle(all_words)
    return all_words, counts


def rebuild_drills_index() -> None:
    """Scan drills/*.html (skip index.html), sort by mtime desc, write index."""
    html_files = [
        p for p in DRILLS_DIR.glob("*.html")
        if p.name != "index.html"
    ]
    if not html_files:
        DRILLS_DIR.joinpath("index.html").write_text(render_drills_placeholder(), encoding="utf-8")
        return

    drill_files = sorted(
        [(p.stem, p.stat().st_mtime) for p in html_files],
        key=lambda x: x[1],
        reverse=True,
    )
    DRILLS_DIR.joinpath("index.html").write_text(
        render_drills_index(drill_files), encoding="utf-8"
    )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Generate a personalized typing drill HTML page from weak bigrams.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "bigrams",
        nargs="*",
        metavar="BIGRAM",
        help="2- or 3-character sequences to drill (e.g. th qu tio)",
    )
    parser.add_argument(
        "--stdin",
        action="store_true",
        help="Read bigrams one per line from stdin instead of positional args",
    )
    parser.add_argument(
        "--name",
        default=None,
        help="Name for the drill (default: weakness-YYYY-MM-DD)",
    )
    parser.add_argument(
        "--corpus",
        default=str(DEFAULT_CORPUS),
        help="Path to corpus wordlist file (default: google-10000-english-usa-no-swears.txt)",
    )
    parser.add_argument(
        "--count",
        type=int,
        default=200,
        help="Total words in output drill (default: 200)",
    )
    parser.add_argument(
        "--out",
        default=None,
        help="Output HTML file path (default: drills/<name>.html)",
    )

    args = parser.parse_args()

    # Collect bigrams
    if args.stdin:
        bigrams = [line.strip() for line in sys.stdin if line.strip()]
    else:
        bigrams = args.bigrams

    if not bigrams:
        print(
            "usage: words-drill [--name NAME] [--corpus PATH] [--count N] [--out PATH] BIGRAM [BIGRAM ...]\n"
            "       words-drill --stdin [other flags]\n"
            "error: no bigrams specified",
            file=sys.stderr,
        )
        sys.exit(2)

    # Resolve name
    today = date.today().isoformat()
    name = args.name if args.name else f"weakness-{today}"

    # Resolve output path
    DRILLS_DIR.mkdir(parents=True, exist_ok=True)
    out_path = Path(args.out) if args.out else DRILLS_DIR / f"{name}.html"

    # Load corpus
    corpus_path = Path(args.corpus)
    corpus = load_corpus(corpus_path)

    # Derive seed from bigrams + name for reproducibility
    seed_str = name + "".join(sorted(bigrams))
    seed = int.from_bytes(seed_str.encode(), "little") % (2**32)

    # Build drill
    words, counts = build_drill_wordlist(corpus, bigrams, args.count, seed)

    # Write HTML
    html = render_drill_page(name, words, bigrams)
    out_path.write_text(html, encoding="utf-8")

    # Rebuild drills index
    rebuild_drills_index()

    # Report
    total_written = len(words)
    detail = ", ".join(f"{b}: {counts.get(b, 0)}" for b in bigrams)
    print(out_path)
    print(
        f"wrote {total_written}/{args.count} words across {len(bigrams)} bigrams ({detail})"
    )


if __name__ == "__main__":
    main()
