"""Generate resource configuration file during build.

This module generates _flash_resource_config.py file with configuration
for ALL resources, which tells the @remote decorator which functions should
execute locally vs create stubs for remote calls based on FLASH_RESOURCE_NAME.
"""

import logging
from pathlib import Path
from typing import Dict, Set

logger = logging.getLogger(__name__)


def generate_all_resource_configs(
    manifest: Dict,
    build_dir: Path,
) -> None:
    """Generate unified resource configuration file for all resources.

    Creates a single _flash_resource_config.py file that contains configuration
    for ALL resources. At runtime, the configuration checks FLASH_RESOURCE_NAME
    to determine which functions are local vs remote for that specific resource.

    Args:
        manifest: Complete Flash manifest with function_registry
        build_dir: Build directory root
    """
    # Get function registry from manifest
    function_registry = manifest.get("function_registry", {})
    resources = manifest.get("resources", {})

    # Build mapping of resource_name -> set of local function names
    resource_functions: Dict[str, Set[str]] = {}

    for resource_name in resources.keys():
        local_funcs = set()
        for func_name, owner_resource in function_registry.items():
            if owner_resource == resource_name:
                local_funcs.add(func_name)
        resource_functions[resource_name] = local_funcs

    # Build set of ALL known functions across all resources
    all_known_functions = set()
    for func_name in function_registry.keys():
        all_known_functions.add(func_name)

    # Generate the unified configuration file
    config_content = '''"""Unified resource configuration for all resources.

THIS FILE IS AUTO-GENERATED BY FLASH BUILD - DO NOT EDIT

Generated configuration tells the @remote decorator which functions
should execute locally vs create stubs for remote calls, based on
the FLASH_RESOURCE_NAME environment variable.
"""

import os
from typing import Set

# Mapping of resource name to set of local function names
RESOURCE_LOCAL_FUNCTIONS = {
'''

    # Add each resource's local functions
    for resource_name, local_funcs in sorted(resource_functions.items()):
        formatted_funcs = _format_set(local_funcs)
        if formatted_funcs == "set()":
            config_content += f'    "{resource_name}": set(),\n'
        else:
            config_content += f'    "{resource_name}": {{{formatted_funcs}}},\n'

    config_content += """}

# Set of all known functions in the manifest
ALL_KNOWN_FUNCTIONS = {"""

    # Add all known functions
    formatted_all_funcs = _format_set(all_known_functions)
    config_content += formatted_all_funcs
    config_content += '''}


def is_local_function(func_name: str) -> bool:
    """Check if a function should execute locally in the current resource.

    Uses FLASH_RESOURCE_NAME environment variable to determine the current
    resource, then checks if the function belongs to that resource.

    Handles the -fb suffix that RunPod backend adds to flashbooted endpoints.

    Args:
        func_name: Name of the function

    Returns:
        True if function should execute locally, False if it needs a stub
    """
    # Get current resource from environment
    current_resource = os.getenv("FLASH_RESOURCE_NAME")

    if not current_resource:
        # No resource name set - assume local for safety
        return True

    # Handle -fb suffix that RunPod backend adds to flashbooted endpoints
    # Try exact match first, then with -fb suffix, then without -fb suffix
    local_functions = RESOURCE_LOCAL_FUNCTIONS.get(current_resource)
    if local_functions is None:
        local_functions = RESOURCE_LOCAL_FUNCTIONS.get(current_resource + "-fb")
    if local_functions is None:
        local_functions = RESOURCE_LOCAL_FUNCTIONS.get(current_resource.removesuffix("-fb"))

    if local_functions is None:
        local_functions = set()

    # Three cases:
    # 1. Function in local set -> execute locally (True)
    # 2. Function known but not in local set -> remote execution (False)
    # 3. Function completely unknown -> execute locally for safety (True)
    if func_name in local_functions:
        return True
    if func_name in ALL_KNOWN_FUNCTIONS:
        return False  # Known to exist in other resource
    return True  # Unknown function, default to local for safety
'''

    # Write the configuration file
    config_file = build_dir / "runpod_flash" / "runtime" / "_flash_resource_config.py"
    config_file.parent.mkdir(parents=True, exist_ok=True)

    with open(config_file, "w", encoding="utf-8") as f:
        f.write(config_content)

    total_functions = sum(len(funcs) for funcs in resource_functions.values())
    logger.debug(
        f"Generated unified resource config: {len(resource_functions)} resources, "
        f"{total_functions} total function mappings"
    )


def _format_set(items: Set[str]) -> str:
    """Format a set of strings for Python code.

    Args:
        items: Set of strings

    Returns:
        Formatted string like '"item1", "item2"' or "set()" for empty set
    """
    if not items:
        return "set()"
    sorted_items = sorted(items)
    return ", ".join(f'"{item}"' for item in sorted_items)
