#!/usr/bin/env python3
"""Search nimi.li for toki pona words by English meaning.

Fetches data from https://nimi.li/__data.json (SvelteKit internal endpoint),
decodes the dedup format, and searches definitions + ku_data keys.

Usage:
    python3 search-nimi.li.py fetch          # Download + cache to cwd
    python3 search-nimi.li.py friend         # Search cached data
    python3 search-nimi.li.py big important  # Multi-word search
"""

import json
import os
import sys
import urllib.request

ENDPOINT = "https://nimi.li/__data.json"
RAW_FILE = "nimi-data-raw.json"
INDEX_FILE = "nimi-words.json"


def fetch_data():
    """Fetch nimi.li data and cache both raw + resolved index to cwd."""
    req = urllib.request.Request(
        ENDPOINT, headers={"User-Agent": "nimi-search/1.0"}
    )
    with urllib.request.urlopen(req) as resp:
        raw = json.loads(resp.read())

    with open(RAW_FILE, "w") as f:
        json.dump(raw, f)

    index = build_index(raw)

    with open(INDEX_FILE, "w") as f:
        json.dump(index, f, indent=2)

    print(f"Fetched {len(index)} words, cached to {RAW_FILE} and {INDEX_FILE}")
    return index


def resolve(pool, val, depth=0):
    """Resolve a value from SvelteKit's shared pool (dedup format).

    The pool is a flat array where dicts/lists/strings are stored once and
    referenced by integer index. This function follows those references
    recursively to produce plain Python objects.
    """
    if depth > 8:
        return val
    if isinstance(val, int) and val < len(pool):
        return resolve(pool, pool[val], depth + 1)
    if isinstance(val, dict):
        return {k: resolve(pool, v, depth + 1) for k, v in val.items()}
    if isinstance(val, list):
        return [resolve(pool, v, depth + 1) for v in val]
    return val


def build_index(raw):
    """Extract word entries from the decoded data.

    Returns a dict mapping word name -> {word, definition, ku_data, ...}.
    """
    pool = raw["nodes"][1]["data"]
    words_dict = pool[1]  # {word_name: index_into_pool}

    index = {}
    for name, word_idx in words_dict.items():
        word_data = resolve(pool, word_idx)
        trans = word_data.get("translations", {})
        index[name] = {
            "word": name,
            "definition": trans.get("definition", ""),
            "commentary": trans.get("commentary", ""),
            "etymology": trans.get("etymology", ""),
            "ku_data": word_data.get("ku_data", {}),
            "source_language": word_data.get("source_language", ""),
            "book": word_data.get("book", ""),
            "usage_category": word_data.get("usage_category", ""),
        }
    return index


def load_index():
    """Load pre-built index from cwd, or fetch if missing."""
    if os.path.exists(INDEX_FILE):
        with open(INDEX_FILE) as f:
            return json.load(f)
    print(f"No cached data found in cwd. Run `{sys.argv[0]} fetch` first.")
    sys.exit(1)


def search(index, query):
    """Search for a word by English meaning.

    Searches definition text, ku_data keys (community usage terms), and
    the word name itself. Returns results sorted by relevance.
    """
    q = query.lower()
    results = []
    for name, data in index.items():
        score = 0
        if q in data["definition"].lower():
            score += 10
        for ku_key in data.get("ku_data", {}):
            if q in ku_key.lower():
                score += 5
                break
        if q in name.lower():
            score += 3
        if q in data.get("commentary", "").lower():
            score += 2
        if score > 0:
            results.append((score, data))
    results.sort(key=lambda x: -x[0])
    return results


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} fetch | <english-word> [english-word ...]")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "fetch":
        fetch_data()
        return

    index = load_index()
    queries = sys.argv[1:]

    for query in queries:
        results = search(index, query)
        if not results:
            print(f"No results for '{query}'.")
            continue
        print(f"\n=== '{query}' ===")
        for _score, r in results:
            print(f"  {r['word']}: {r['definition']}")
            if r["commentary"]:
                print(f"    ({r['commentary'][:120]})")


if __name__ == "__main__":
    main()