#!/usr/bin/env python3
"""Expand kb.source.json into matchable terms and emit the JS + Swift knowledge bases.

One source, two targets, so the website and the rNitro app can never drift apart.
Run with --stats for counts, --check to validate without writing.
"""
from __future__ import annotations

import json
import re
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
SOURCE = HERE / "kb.source.json"
JS_OUT = HERE.parent / "chopstickshq-site" / "js" / "chopsticks-ai-kb.js"
SWIFT_OUT = HERE / "ChopsticksAIKB.swift"
JSON_OUT = HERE.parent / "chopstickshq-site" / "api" / "_lib" / "chopsticks-ai-kb.json"

# Published copies shown on /chopsticks-ai/. Refreshed here so they cannot drift
# from the real sources the way a manual cp inevitably does.
FILES_OUT = HERE.parent / "chopstickshq-site" / "chopsticks-ai" / "files"
PUBLISHED = [
    (HERE / "kb.source.json", "kb.source.json"),
    (HERE / "build-kb.py", "build-kb.py"),
    (HERE / "ChopsticksAIEngine.swift", "ChopsticksAIEngine.swift"),
    (HERE / "ChopsticksAIKB.swift", "ChopsticksAIKB.swift"),
    (HERE / "fixtures.json", "fixtures.json"),
    (HERE / "run-fixtures.mjs", "run-fixtures.mjs"),
    (JS_OUT, "chopsticks-ai-kb.js"),
    (HERE.parent / "chopstickshq-site" / "js" / "chopsticks-ai.js", "chopsticks-ai.web.js"),
    (HERE.parent / "chopstickshq-site" / "api" / "_lib" / "chopsticks-ai.js", "chopsticks-ai.server.js"),
]

# Words too generic to identify an intent on their own. A term that reduces to
# only these is dropped rather than allowed to match everything.
STOPWORDS = {
    "a", "an", "and", "are", "at", "be", "by", "can", "did", "do", "does", "for",
    "from", "get", "got", "has", "have", "how", "i", "if", "in", "is", "it", "its",
    "me", "my", "no", "not", "of", "on", "or", "so", "the", "then", "there", "this",
    "to", "up", "use", "was", "what", "when", "where", "which", "who", "why", "will",
    "with", "you", "your",
}

MIN_TERM_LEN = 3

# Adjacent keys on a QWERTY board, used to synthesise realistic typos.
ADJACENT = {
    "a": "qws", "b": "vgn", "c": "xdv", "d": "serfcx", "e": "wsdr", "f": "drtgvc",
    "g": "ftyhbv", "h": "gyujnb", "i": "ujko", "j": "huikmn", "k": "jiolm",
    "l": "kop", "m": "njk", "n": "bhjm", "o": "iklp", "p": "ol", "q": "wa",
    "r": "edft", "s": "awedxz", "t": "rfgy", "u": "yhji", "v": "cfgb",
    "w": "qase", "x": "zsdc", "y": "tghu", "z": "asx",
}


def normalise(text: str) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9\s]+", " ", text)
    return re.sub(r"\s+", " ", text).strip()


# Endings that are verbs or participles rather than nouns. Pluralising these
# produces junk like "blockeds" or "verifys", which no one would ever type.
NO_PLURAL_SUFFIX = ("ed", "ing", "ly", "s", "fy", "y")


def morphological(term: str) -> set[str]:
    """Plural/singular and separator variants: menu bar / menubar / menu-bar."""
    out = {term}
    words = term.split()

    # Only compound short two-word terms. Squashing a whole sentence gives
    # "applecouldnotverify", which is noise nobody types.
    if len(words) == 2 and len(term) <= 14:
        out.add(term.replace(" ", ""))
        out.add(term.replace(" ", "-"))

    # Pluralise single nouns only. Inflecting the last word of a phrase turns
    # "apple could not verify" into "apple could not verifys".
    if len(words) == 1:
        w = words[0]
        if len(w) >= 4 and w.isalpha():
            if w.endswith("s") and not w.endswith("ss"):
                out.add(w[:-1])
            elif not w.endswith(NO_PLURAL_SUFFIX):
                out.add(w + "s")
    return out


def typos(word: str, limit: int = 6) -> set[str]:
    """Transposition and adjacent-key typos for a single word."""
    out: set[str] = set()
    if len(word) < 5 or not word.isalpha():
        return out
    for i in range(len(word) - 1):
        if word[i] != word[i + 1]:
            out.add(word[:i] + word[i + 1] + word[i] + word[i + 2:])
        if len(out) >= limit:
            return set(sorted(out)[:limit])
    for i, ch in enumerate(word):
        for repl in ADJACENT.get(ch, ""):
            out.add(word[:i] + repl + word[i + 1:])
            if len(out) >= limit:
                return set(sorted(out)[:limit])
    return set(sorted(out)[:limit])


def build_synonym_map(raw: dict[str, list[str]]) -> dict[str, set[str]]:
    """Bidirectional: every member of a group expands to every other member."""
    groups: dict[str, set[str]] = {}
    for head, members in raw.items():
        group = {normalise(head)} | {normalise(m) for m in members}
        for member in group:
            groups.setdefault(member, set()).update(group)
    return groups


def expand(term: str, synonyms: dict[str, set[str]]) -> set[str]:
    base = normalise(term)
    if not base:
        return set()

    seeded = {base} | synonyms.get(base, set())

    # Substitute a synonym for any single word inside a multi-word term.
    words = base.split()
    if len(words) > 1:
        for i, w in enumerate(words):
            for alt in synonyms.get(w, set()):
                seeded.add(" ".join(words[:i] + [alt] + words[i + 1:]))

    forms: set[str] = set()
    for s in seeded:
        forms |= morphological(s)

    # Typos only for genuinely single-word terms. Generating them from squashed
    # compounds produces unusable strings like "aplpecouldnotverify".
    single_words = {s for s in seeded if " " not in s}
    for f in list(forms):
        if " " not in f and "-" not in f and f in single_words:
            forms |= typos(f)

    cleaned = set()
    for f in forms:
        f = normalise(f)
        if len(f) < MIN_TERM_LEN:
            continue
        parts = f.split()
        # A lone stopword matches everything, so drop it. A *phrase* built only
        # from stopwords ("who are you", "what can you do") is still a specific
        # question, so it survives.
        if all(w in STOPWORDS for w in parts) and len(parts) < 3:
            continue
        cleaned.add(f)
    return cleaned


def weight_for(term: str, kind: str) -> int:
    """Phrases beat multi-word keywords beat single tokens; length breaks ties."""
    words = len(term.split())
    if kind == "phrase":
        base = 6
    elif words > 1:
        base = 4
    else:
        base = 2
    return base + min(len(term) // 6, 2)


def check_collisions(intents: list[dict]) -> list[str]:
    """Authored-term collisions between same-priority intents are ambiguous."""
    errors: list[str] = []

    seen_ids: set[str] = set()
    for it in intents:
        if it["id"] in seen_ids:
            errors.append(f"duplicate intent id: {it['id']}")
        seen_ids.add(it["id"])

    owners: dict[str, list[dict]] = {}
    for it in intents:
        for term in it.get("keywords", []) + it.get("phrases", []):
            owners.setdefault(normalise(term), []).append(it)

    for term, claimers in sorted(owners.items()):
        if len(claimers) < 2:
            continue
        priorities = {c["priority"] for c in claimers}
        if len(priorities) < len(claimers):
            ids = ", ".join(sorted(c["id"] for c in claimers))
            errors.append(
                f"ambiguous keyword {term!r} claimed by same-priority intents: {ids}"
            )
    return errors


def build() -> tuple[dict, dict]:
    data = json.loads(SOURCE.read_text())
    synonyms = build_synonym_map(data.get("synonyms", {}))
    intents = data["intents"]

    errors = check_collisions(intents)
    if errors:
        for e in errors:
            print(f"ERROR: {e}", file=sys.stderr)
        raise SystemExit(1)

    compiled = []
    total_terms = 0
    for it in intents:
        terms: dict[str, int] = {}
        for kind, field in (("phrase", "phrases"), ("keyword", "keywords")):
            for raw in it.get(field, []):
                for term in expand(raw, synonyms):
                    w = weight_for(term, kind)
                    if w > terms.get(term, 0):
                        terms[term] = w
        total_terms += len(terms)
        compiled.append({
            "id": it["id"],
            "product": it["product"],
            "label": it["label"],
            "priority": it["priority"],
            "answer": it["answer"],
            # sorted for deterministic output
            "terms": [[t, terms[t]] for t in sorted(terms)],
        })

    compiled.sort(key=lambda c: c["id"])
    kb = {"meta": data["meta"], "intents": compiled}
    stats = {
        "intents": len(compiled),
        "authored": sum(
            len(i.get("keywords", [])) + len(i.get("phrases", [])) for i in intents
        ),
        "expanded": total_terms,
        "products": sorted({i["product"] for i in compiled}),
    }
    return kb, stats


def emit_js(kb: dict) -> str:
    payload = json.dumps(kb, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    return (
        "/* Generated by chopsticks-ai/build-kb.py - do not edit by hand. */\n"
        f"window.CHOPSTICKS_AI_KB = {payload};\n"
    )


def swift_string(s: str) -> str:
    return json.dumps(s, ensure_ascii=False)


def emit_swift(kb: dict) -> str:
    lines = [
        "// Generated by chopsticks-ai/build-kb.py - do not edit by hand.",
        "",
        "struct ChopsticksAIIntent {",
        "    let id: String",
        "    let product: String",
        "    let label: String",
        "    let priority: Int",
        "    let answer: String",
        "    let terms: [(String, Int)]",
        "}",
        "",
        "enum ChopsticksAIKB {",
        f"    static let tagline = {swift_string(kb['meta']['tagline'])}",
        f"    static let version = {swift_string(kb['meta']['version'])}",
        "    static let intents: [ChopsticksAIIntent] = [",
    ]
    for it in kb["intents"]:
        terms = ", ".join(f"({swift_string(t)}, {w})" for t, w in it["terms"])
        lines += [
            "        ChopsticksAIIntent(",
            f"            id: {swift_string(it['id'])},",
            f"            product: {swift_string(it['product'])},",
            f"            label: {swift_string(it['label'])},",
            f"            priority: {it['priority']},",
            f"            answer: {swift_string(it['answer'])},",
            f"            terms: [{terms}]",
            "        ),",
        ]
    lines += ["    ]", "}", ""]
    return "\n".join(lines)


def main() -> None:
    args = set(sys.argv[1:])
    kb, stats = build()

    if "--stats" in args or "--check" in args:
        print(f"intents        : {stats['intents']}")
        print(f"authored terms : {stats['authored']}")
        print(f"expanded terms : {stats['expanded']}")
        print(f"products       : {', '.join(stats['products'])}")

    if "--check" in args:
        print("check passed - nothing written")
        return

    JS_OUT.parent.mkdir(parents=True, exist_ok=True)
    JS_OUT.write_text(emit_js(kb))
    SWIFT_OUT.write_text(emit_swift(kb))
    JSON_OUT.parent.mkdir(parents=True, exist_ok=True)
    JSON_OUT.write_text(json.dumps(kb, separators=(",", ":"), sort_keys=True))
    print(f"wrote {JS_OUT.relative_to(HERE.parent)} ({JS_OUT.stat().st_size:,} bytes)")
    print(f"wrote {SWIFT_OUT.relative_to(HERE.parent)} ({SWIFT_OUT.stat().st_size:,} bytes)")
    print(f"wrote {JSON_OUT.relative_to(HERE.parent)} ({JSON_OUT.stat().st_size:,} bytes)")

    FILES_OUT.mkdir(parents=True, exist_ok=True)
    n = 0
    for src, name in PUBLISHED:
        if src.is_file():
            (FILES_OUT / name).write_bytes(src.read_bytes())
            n += 1
    print(f"refreshed {n} published source files in {FILES_OUT.relative_to(HERE.parent)}/")


if __name__ == "__main__":
    main()