import json
import os
import sys
from pathlib import Path

LOCALES_DIR = Path("src/locales")
USER_MASTER = LOCALES_DIR / "translations.user.json"
ADMIN_MASTER = LOCALES_DIR / "translations.admin.json"
LANGUAGES = ["en", "me", "ru", "ua"]

def get_nested_keys(data, prefix=""):
    keys = {}
    for k, v in data.items():
        new_prefix = f"{prefix}.{k}" if prefix else k
        if isinstance(v, dict):
            if any(lang in v for lang in LANGUAGES):
                keys[new_prefix] = v
            else:
                keys.update(get_nested_keys(v, new_prefix))
        else:
            keys[new_prefix] = v
    return keys

def set_nested_key(data, key_path, value):
    parts = key_path.split('.')
    for part in parts[:-1]:
        data = data.setdefault(part, {})
    data[parts[-1]] = value

def deep_merge(dict1, dict2):
    """Recursive merge of two dictionaries."""
    for key, value in dict2.items():
        if key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict):
            # If both are dicts and it's NOT a leaf node with languages
            if not any(lang in value for lang in LANGUAGES):
                deep_merge(dict1[key], value)
            else:
                dict1[key] = value
        else:
            dict1[key] = value

def assemble_master(fragment_dir, master_file, root_key=None):
    """Merge all fragments in a directory into one master file"""
    if not fragment_dir.exists():
        print(f"Warning: {fragment_dir} not found.")
        return
    
    combined_data = {}
    # Sort files to ensure deterministic merging
    for item in sorted(fragment_dir.glob("*.json")):
        with open(item, "r", encoding="utf-8") as f:
            try:
                data = json.load(f)
                deep_merge(combined_data, data)
            except Exception as e:
                print(f"Error reading {item}: {e}")

    # Wrap in root key if needed (e.g. "admin": {...})
    final_output = combined_data
    if root_key:
        final_output = {root_key: combined_data}

    with open(master_file, "w", encoding="utf-8") as f:
        json.dump(final_output, f, ensure_ascii=False, indent=2)
    print(f"Assembled fragments from {fragment_dir} into {master_file}")

def _merge_files(master_file, suffix):
    """Generic merge from {lang}{suffix} into master_file"""
    master_data = {}
    all_keys = set()
    locale_data = {}

    for lang in LANGUAGES:
        path = LOCALES_DIR / f"{lang}{suffix}"
        if path.exists():
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
                flat = {}
                def flatten(d, prefix=""):
                    for k, v in d.items():
                        new_prefix = f"{prefix}.{k}" if prefix else k
                        if isinstance(v, dict):
                            flatten(v, new_prefix)
                        else:
                            flat[new_prefix] = v
                flatten(data)
                locale_data[lang] = flat
                all_keys.update(flat.keys())

    for key in sorted(all_keys):
        translations = {}
        for lang in LANGUAGES:
            translations[lang] = locale_data.get(lang, {}).get(key, "")
        set_nested_key(master_data, key, translations)

    with open(master_file, "w", encoding="utf-8") as f:
        json.dump(master_data, f, ensure_ascii=False, indent=2)
    print(f"Merged {suffix} files into {master_file}")

def _split_master(master_file, suffix):
    """Generic split from master_file into {lang}{suffix}"""
    if not master_file.exists():
        print(f"Warning: {master_file} not found, skipping split.")
        return

    with open(master_file, "r", encoding="utf-8") as f:
        master_data = json.load(f)

    for lang in LANGUAGES:
        lang_data = {}
        def unflatten(d, target, current_lang):
            for k, v in d.items():
                if isinstance(v, dict):
                    if any(l in v for l in LANGUAGES):
                        target[k] = v.get(current_lang, "")
                    else:
                        target[k] = {}
                        unflatten(v, target[k], current_lang)
                else:
                    target[k] = v
        
        unflatten(master_data, lang_data, lang)
        output_path = LOCALES_DIR / f"{lang}{suffix}"
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(lang_data, f, ensure_ascii=False, indent=2)
        print(f"Generated {output_path}")

def assemble():
    """Assemble fragments into masters"""
    assemble_master(LOCALES_DIR / "master_user", USER_MASTER)
    assemble_master(LOCALES_DIR / "master_admin", ADMIN_MASTER, root_key="admin")

def merge():
    """Merge individual JSON files into respective master files"""
    _merge_files(USER_MASTER, ".json")
    _merge_files(ADMIN_MASTER, ".admin.json")

def split():
    """Assemble and then split master files"""
    assemble()
    _split_master(USER_MASTER, ".json")
    _split_master(ADMIN_MASTER, ".admin.json")

def list_missing():
    """List all keys with missing translations in both master files"""
    for name, path in [("User", USER_MASTER), ("Admin", ADMIN_MASTER)]:
        if not path.exists():
            continue
        print(f"\n--- Missing in {name} ---")
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        flat_keys = get_nested_keys(data)
        missing_found = False
        for key, trans in sorted(flat_keys.items()):
            if isinstance(trans, dict):
                missing = [l for l in LANGUAGES if not trans.get(l)]
                if missing:
                    print(f"Key: {key} - Missing: {', '.join(missing)}")
                    missing_found = True
        if not missing_found:
            print("Everything translated!")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python manage_locales.py [assemble|merge|split|missing]")
        sys.exit(1)
    
    cmd = sys.argv[1].lower()
    if cmd == "assemble":
        assemble()
    elif cmd == "merge":
        merge()
    elif cmd == "split":
        split()
    elif cmd == "missing":
        list_missing()
    else:
        print(f"Unknown command: {cmd}")