import os
import json
from collections import defaultdict
from utils.unicleaner import clean_unicode

BRAINMAP_FILE = "data/memory/brainmap.json"


def load_brainmap():
    if os.path.exists(BRAINMAP_FILE):
        with open(BRAINMAP_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}


def save_brainmap(map_data):
    with open(BRAINMAP_FILE, "w", encoding="utf-8") as f:
        json.dump(map_data, f, indent=2)


brain_map = load_brainmap()


def update_brainmap(words):
    for i, word in enumerate(words):
        for j in range(i+1, len(words)):
            w1 = word
            w2 = words[j]
            if w1 == w2:
                continue
            if w1 not in brain_map:
                brain_map[w1] = {}
            if w2 not in brain_map[w1]:
                brain_map[w1][w2] = 0
            brain_map[w1][w2] += 1
    save_brainmap(brain_map)


def get_brainmap():
    return brain_map


def fix_brainmap(brainmap: dict) -> dict:
    cleaned_brainmap = {}

    for word, value in brainmap.items():
        cleaned_word = clean_unicode(word.strip())

        # Skip bad entries
        if not cleaned_word or cleaned_word in {"...", "-", "--", "''", '""'}:
            continue

        # Merge duplicates (case-insensitive optional)
        if cleaned_word in cleaned_brainmap:
            cleaned_brainmap[cleaned_word] += value
        else:
            cleaned_brainmap[cleaned_word] = value

    return cleaned_brainmap