"""
Orchestration logic for RECodeGenerator.
"""
from __future__ import annotations
from pathlib import Path
from typing import Any, Dict, List
from .models import LinkSpec, IndependentObjects
from .utils import safe_name


def create_folder_structure(folder_location: Path, groups: List[Any], independent: IndependentObjects) -> None:
    """Create the folder structure with groups/ and independent/ subdirectories."""
    folder_location.mkdir(parents=True, exist_ok=True)
    
    # Create groups/ subdirectory
    groups_folder = folder_location / "groups"
    groups_folder.mkdir(exist_ok=True)
    
    # Create independent/ subdirectory
    independent_folder = folder_location / "independent"
    independent_folder.mkdir(exist_ok=True)
    
    # Generate main __init__.py
    init_lines = [
        '"""',
        'Generated by RECodeGenerator',
        '',
        f'Contains {len(groups)} groups with units, graphics, and sounds.',
        '"""',
        'import sys',
        'from pathlib import Path',
        '',
        '# Add this directory to path for imports',
        '_pkg_dir = Path(__file__).parent',
        'if str(_pkg_dir) not in sys.path:',
        '    sys.path.insert(0, str(_pkg_dir))',
        '',
        '# Import all apply functions for easy access',
    ]
    
    for group in groups:
        safe_grp_name = safe_name(group.name)
        init_lines.append(
            f'from groups.{group.folder_name}.apply_{safe_grp_name} import apply_{safe_grp_name}_group'
        )
    
    init_lines.extend([
        '',
        '# Import master module',
        'from reverse_engineer_dat import reverse_engineer_dat, list_groups',
        '',
    ])
    
    (folder_location / "__init__.py").write_text("\n".join(init_lines))
    
    # Create groups/__init__.py
    groups_init_lines = [
        '"""Groups containing related units, graphics, and sounds."""',
    ]
    (groups_folder / "__init__.py").write_text("\n".join(groups_init_lines))
    
    # Create group folders inside groups/
    for group in groups:
        group_folder = groups_folder / group.folder_name
        group_folder.mkdir(exist_ok=True)
        
        safe_grp_name = safe_name(group.name)
        group_init_lines = [
            f'"""Group: {group.name}"""',
            'import sys',
            'from pathlib import Path',
            '',
            '_group_dir = Path(__file__).parent',
            'if str(_group_dir) not in sys.path:',
            '    sys.path.insert(0, str(_group_dir))',
            '',
            f'from apply_{safe_grp_name} import apply_{safe_grp_name}_group',
            'from units import *',
            'from graphics import *', 
            'from sounds import *',
            '',
        ]
        
        (group_folder / "__init__.py").write_text("\n".join(group_init_lines))
    
    # Create independent/__init__.py
    independent_init_lines = [
        '"""Independent objects with no links to other above-threshold objects."""',
        'import sys',
        'from pathlib import Path',
        '',
        '_indep_dir = Path(__file__).parent',
        'if str(_indep_dir) not in sys.path:',
        '    sys.path.insert(0, str(_indep_dir))',
        '',
        'from units import *',
        'from graphics import *',
        'from sounds import *',
        '',
    ]
    (independent_folder / "__init__.py").write_text("\n".join(independent_init_lines))


def generate_apply_module(ws: Any, group: Any, folder_location: Path,
                         all_unit_links: Dict[int, List[LinkSpec]],
                         all_graphic_links: Dict[int, List[LinkSpec]],
                         all_sound_links: Dict[int, List[LinkSpec]]) -> None:
    """Generate the apply_group_X.py module."""
    group_folder = folder_location / "groups" / group.folder_name
    
    lines = [
        '"""',
        f'Apply Group: {group.name}',
        f'',
        f'Creates all objects and links them together.',
        '"""',
        'import sys',
        'import importlib.util',
        'from pathlib import Path',
        '',
        '# Load local modules',
        '_current_dir = Path(__file__).parent',
        '',
        'def _load_local_module(name):',
        '    spec = importlib.util.spec_from_file_location(name, _current_dir / f"{name}.py")',
        '    mod = importlib.util.module_from_spec(spec)',
        '    spec.loader.exec_module(mod)',
        '    return mod',
        '',
        '_units = _load_local_module("units")',
        '_graphics = _load_local_module("graphics")',
        '_sounds = _load_local_module("sounds")',
        '',
        '',
        f'def apply_{safe_name(group.name)}_group(ws):',
        f'    """Create and link all objects in {group.name} group."""',
        f'    print(f"  > Applying Group: {group.name}")',
        '',
    ]
    
    # 1. Create Graphics
    lines.append("    # Create Graphics")
    graphic_vars = {}
    for gid in sorted(group.graphic_ids):
        try:
            graphic = ws.graphic_manager.get(gid)
            name = getattr(graphic, "name", f"Graphic_{gid}")
        except Exception:
            name = f"Graphic_{gid}"
        
        gfx_safe_name = safe_name(name)
        var_name = f"graphic_{gid}"
        graphic_vars[gid] = var_name
        lines.append(f"    {var_name} = _graphics.create_graphic_{gid}_{gfx_safe_name}(ws)")
        
    lines.append("")
    
    # 2. Create Sounds
    lines.append("    # Create Sounds")
    sound_vars = {}
    for sid in sorted(group.sound_ids):
        try:
            sound = ws.sound_manager.get(sid)
            sound_name = f"sound_{sid}"
            sound_files = sound.sounds if hasattr(sound, 'sounds') else []
            if sound_files:
                first_file = sound_files[0]
                if hasattr(first_file, 'sound_name') and first_file.sound_name:
                    sound_name = first_file.sound_name
                elif hasattr(first_file, 'filename') and first_file.filename:
                    sound_name = first_file.filename
        except Exception:
            sound_name = f"sound_{sid}"

        snd_safe_name = safe_name(sound_name)
        var_name = f"sound_{sid}"
        sound_vars[sid] = var_name
        lines.append(f"    {var_name} = _sounds.create_sound_{sid}_{snd_safe_name}(ws)")
        
    lines.append("")
    
    # 3. Create Units
    lines.append("    # Create Units")
    unit_vars = {}
    for uid in sorted(group.unit_ids):
        try:
            unit = ws.unit_manager.get(uid)
            name = getattr(unit, "name", f"Unit_{uid}")
        except Exception:
            name = f"Unit_{uid}"
        
        unit_safe_name = safe_name(name)
        var_name = f"unit_{uid}"
        unit_vars[uid] = var_name
        lines.append(f"    {var_name} = _units.create_unit_{uid}_{unit_safe_name}(ws)")
        
    lines.append("")
    lines.append("    # Apply Links")
    
    # Link Graphics
    for gid in sorted(group.graphic_ids):
        links = all_graphic_links.get(gid, [])
        if not links:
            continue

        var_name = graphic_vars.get(gid, f"graphic_{gid}")
        for link in sorted(links, key=lambda l: l.field):
            setter = f"{var_name}.{link.field}"

            if link.target_type == "graphic" and link.target_id in graphic_vars:
                lines.append(f"    {setter} = {graphic_vars[link.target_id]}.id")
            elif link.target_type == "sound" and link.target_id in sound_vars:
                lines.append(f"    {setter} = {sound_vars[link.target_id]}.id")
            else:
                lines.append(f"    {setter} = {link.target_id}  # external {link.target_type}")

    # Link Sounds
    for sid in sorted(group.sound_ids):
        links = all_sound_links.get(sid, [])
        if not links:
            continue

        var_name = sound_vars.get(sid, f"sound_{sid}")
        for link in sorted(links, key=lambda l: l.field):
            setter = f"{var_name}.{link.field}"
            lines.append(f"    {setter} = {link.target_id}")

    # Link Units
    for uid in sorted(group.unit_ids):
        links = all_unit_links.get(uid, [])
        if not links:
            continue

        var_name = unit_vars.get(uid, f"unit_{uid}")

        for link in sorted(links, key=lambda l: l.field):
            setter = f"{var_name}.{link.field}"

            if link.target_type == "unit" and link.target_id in unit_vars:
                lines.append(f"    {setter} = {unit_vars[link.target_id]}.id")
            elif link.target_type == "graphic" and link.target_id in graphic_vars:
                lines.append(f"    {setter} = {graphic_vars[link.target_id]}.id")
            elif link.target_type == "sound" and link.target_id in sound_vars:
                lines.append(f"    {setter} = {sound_vars[link.target_id]}.id")
            else:
                lines.append(f"    {setter} = {link.target_id}  # external {link.target_type}")
    
    lines.append("")
    lines.append("    # Return main unit")
    if unit_vars:
        first_var = list(unit_vars.values())[0]
        lines.append(f"    return {first_var}")
    else:
        lines.append("    return None")
    lines.append("")
    
    apply_path = group_folder / f"apply_{safe_name(group.name)}.py"
    apply_path.write_text("\n".join(lines))


def generate_master_module(folder_location: Path, groups: List[Any], independent: IndependentObjects) -> None:
    """Generate the master reverse_engineer_dat.py module."""
    lines = [
        '"""',
        'Reverse Engineer DAT - Master Apply Module',
        '',
        f'Generated {len(groups)} groups',
        'Master Module: reverse_engineer_dat',
        '',
        'Orchestrates the reverse engineering of all grouped units.',
        '"""',
        'import sys',
        'import os',
        'from pathlib import Path',
        '',
        '# Add current directory to sys.path',
        'base_dir = Path(__file__).resolve().parent',
        'if str(base_dir) not in sys.path:',
        '    sys.path.insert(0, str(base_dir))',
        '',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    # Import group apply functions
    for group in groups:
        safe_grp_name = safe_name(group.name)
        lines.extend([
            f'try:',
            f'    from groups.{group.folder_name}.apply_{safe_grp_name} import apply_{safe_grp_name}_group',
            f'except ImportError:',
            f'    from .groups.{group.folder_name}.apply_{safe_grp_name} import apply_{safe_grp_name}_group'
        ])
    
    # Import independent objects
    lines.extend([
        '',
        '# Import independent objects',
        'try:',
        '    from independent.units import *',
        '    from independent.graphics import *',
        '    from independent.sounds import *',
        'except ImportError:',
        '    from .independent.units import *',
        '    from .independent.graphics import *',
        '    from .independent.sounds import *',
    ])
    
    lines.extend([
        '',
        '',
        'def reverse_engineer_dat(ws):',
        '    """Apply all groups and independent objects to recreate units."""',
        '    results = {}',
        f'    total_groups = {len(groups)}',
        '    ',
    ])
    
    # Apply independent objects first
    lines.append('    # Apply independent objects')
    lines.append('    print("Applying independent objects...")')
    # Note: We'll add specific calls when generating the master - this is done at generate time
    
    lines.append('')
    lines.append('    # Apply groups')
    for i, group in enumerate(groups):
        safe_grp_name = safe_name(group.name)
        lines.append(f'    print(f"[{i+1}/{{total_groups}}] Processing group: {group.name}...")')
        lines.append(f'    results["{safe_grp_name}"] = apply_{safe_grp_name}_group(ws)')
    
    lines.extend([
        '',
        '    return results',
        '',
        'def list_groups():',
        '    """List names of all generated groups."""',
        '    return [',
    ])
    
    for group in groups:
        lines.append(f'        "{group.name}",')
        
    lines.extend([
        '    ]',
        '',
        '',
        'if __name__ == "__main__":',
        '    import argparse',
        '    ',
        '    parser = argparse.ArgumentParser(description="Reverse engineer DAT file")',
        '    parser.add_argument("--input", "-i", default="empires2_x2_p1.dat", help="Input DAT file")',
        '    parser.add_argument("--output", "-o", default="Rebuild_Output.dat", help="Output DAT file")',
        '    args = parser.parse_args()',
        '    ',
        '    print(f"Loading workspace from {args.input}...")',
        '    ws = GenieWorkspace.load(args.input)',
        '    ',
        '    print(f"Applying {len(list_groups())} groups...")',
        '    results = reverse_engineer_dat(ws)',
        '    ',
        '    print(f"Saving to {args.output}...")',
        '    ws.save(args.output, validate=False)',
        '    print("Done!")',
    ])
    
    (folder_location / "reverse_engineer_dat.py").write_text("\n".join(lines))


def generate_group_code(ws: Any, group: Any, folder_location: Path, config: Any) -> None:
    """Generate all code for a group."""
    from .writers import generate_unit_code, generate_graphic_code, generate_sound_code
    
    group_folder = folder_location / "groups" / group.folder_name
    
    # Track all links
    all_unit_links: Dict[int, List[LinkSpec]] = {}
    all_graphic_links: Dict[int, List[LinkSpec]] = {}
    all_sound_links: Dict[int, List[LinkSpec]] = {}
    
    # ===== UNITS =====
    units_lines = [
        '"""',
        f'Units for group: {group.name}',
        f'Unit IDs: {sorted(group.unit_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for unit_id in sorted(group.unit_ids):
        unit_code, links = generate_unit_code(ws, unit_id, group, config)
        units_lines.extend(unit_code)
        all_unit_links[unit_id] = links
    
    (group_folder / "units.py").write_text("\n".join(units_lines))
    
    # ===== GRAPHICS =====
    graphics_lines = [
        '"""',
        f'Graphics for group: {group.name}',
        f'Graphic IDs: {sorted(group.graphic_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for gid in sorted(group.graphic_ids):
        gfx_code, links = generate_graphic_code(ws, gid, group, config)
        graphics_lines.extend(gfx_code)
        all_graphic_links[gid] = links
    
    (group_folder / "graphics.py").write_text("\n".join(graphics_lines))
    
    # ===== SOUNDS =====
    sounds_lines = [
        '"""',
        f'Sounds for group: {group.name}',
        f'Sound IDs: {sorted(group.sound_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for sid in sorted(group.sound_ids):
        snd_code, links = generate_sound_code(ws, sid, group, config)
        sounds_lines.extend(snd_code)
        all_sound_links[sid] = links
    
    (group_folder / "sounds.py").write_text("\n".join(sounds_lines))
    
    # ===== APPLY MODULE =====
    generate_apply_module(ws, group, folder_location, all_unit_links, all_graphic_links, all_sound_links)


def generate_independent_code(ws: Any, independent: IndependentObjects, folder_location: Path, config: Any) -> None:
    """Generate code for independent objects."""
    from .writers import (
        generate_independent_unit_code,
        generate_independent_graphic_code,
        generate_independent_sound_code,
    )
    
    independent_folder = folder_location / "independent"
    
    # ===== UNITS =====
    units_lines = [
        '"""',
        'Independent Units',
        f'Unit IDs: {sorted(independent.unit_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for unit_id in sorted(independent.unit_ids):
        unit_code = generate_independent_unit_code(ws, unit_id, config)
        units_lines.extend(unit_code)
    
    (independent_folder / "units.py").write_text("\n".join(units_lines))
    
    # ===== GRAPHICS =====
    graphics_lines = [
        '"""',
        'Independent Graphics',
        f'Graphic IDs: {sorted(independent.graphic_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for gid in sorted(independent.graphic_ids):
        gfx_code = generate_independent_graphic_code(ws, gid, config)
        graphics_lines.extend(gfx_code)
    
    (independent_folder / "graphics.py").write_text("\n".join(graphics_lines))
    
    # ===== SOUNDS =====
    sounds_lines = [
        '"""',
        'Independent Sounds',
        f'Sound IDs: {sorted(independent.sound_ids)}',
        '"""',
        'from aoe2_genie_tooling.Base.workspace import GenieWorkspace',
        '',
    ]
    
    for sid in sorted(independent.sound_ids):
        snd_code = generate_independent_sound_code(ws, sid, config)
        sounds_lines.extend(snd_code)
    
    (independent_folder / "sounds.py").write_text("\n".join(sounds_lines))
