import csv
import gzip
import json
import logging
from collections import defaultdict
from pathlib import Path

from zangramru.utils.trie import TrieNode

CWD = Path(__file__).parent
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def build_trie(dict_path: Path) -> TrieNode:
    """Build trie for wikidict definitions."""
    trie = TrieNode()
    total_words = 0
    with gzip.open(dict_path) as f_in:
        for line in f_in:
            word_info = json.loads(line)
            # If it starts wit capital.
            if word_info.get("lang_code") != "ru":
                continue
            if word_info.get("pos") != "noun":
                continue
            if word_info["word"][0].isupper():
                continue
            senses = word_info.get("senses", [])
            multi_def = []
            for sense in senses:
                glosses = sense.get("glosses")
                if glosses:
                    multi_def.extend(glosses)

            total_words += 1
            trie.add(word_info["word"], definitions=multi_def)

    logger.info(f"Trie is ready. Total words: {total_words}")
    return trie


def get_defs(
    dict_file: Path,
    defs: TrieNode,
    stats: dict[str, int],
) -> dict[str, list[str]]:
    """Find all word definitions for all our dicts."""
    word_def = {}
    with dict_file.open("r") as f_in:
        for line in f_in:
            if len(line.strip()) < 4:
                stats["too_short"] += 1
                continue
            node = defs.find_node(line)
            if node is None:
                stats["not_found"] += 1
                continue
            if not node.valid_word:
                stats["invalid"] += 1
                continue
            word_def[node.value] = node.definitions
    return word_def


def main() -> None:
    """Update dictionaries."""
    trie = build_trie(Path(CWD / "wikidict.jsonl.gz"))
    with Path(CWD.parent.parent / "zangramru/services/dict.csv").open("w") as f_out:
        writer = csv.writer(f_out, delimiter="|")
        writer.writerow(["word", "def"])
        stats = defaultdict(int)
        words: dict[str, list[str]] = {}

        for file in Path(CWD / "data").iterdir():
            logger.info(f"Processing {file}")
            words.update(get_defs(file, trie, stats))

        logger.info("Done updating")
        logger.info(stats)

        for word, defs in words.items():
            all_defs = ""
            if len(defs) == 1:
                all_defs = defs[0]
            else:
                all_defs = "\n".join(f"{i + 1}: {defn}" for i, defn in enumerate(defs))
            writer.writerow([word, all_defs])


if __name__ == "__main__":
    main()