"""Generator for handler_<name>.py files."""

import ast
import logging
from pathlib import Path
from typing import Any, Dict, List, Union

from runpod_flash.runtime.models import Manifest

logger = logging.getLogger(__name__)

HANDLER_TEMPLATE = '''"""
Auto-generated handler for resource: {resource_name}
Generated at: {timestamp}

This file is generated by the Flash build process. Do not edit manually.
"""

import importlib
from runpod_flash.runtime.generic_handler import create_handler

# Import all functions/classes that belong to this resource
{imports}

# Function registry for this handler
FUNCTION_REGISTRY = {{
{registry}
}}

# Create configured handler
handler = create_handler(FUNCTION_REGISTRY)

if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{"handler": handler}})
'''

DEPLOYED_HANDLER_TEMPLATE = '''"""
Auto-generated deployed handler for resource: {resource_name}
Generated at: {timestamp}

Deployed endpoint handler: accepts plain JSON, no cloudpickle.
One function per endpoint, identified by FLASH_RESOURCE_NAME.

This file is generated by the Flash build process. Do not edit manually.
"""

import asyncio
import importlib
import inspect
import logging
import traceback

_logger = logging.getLogger(__name__)

# Import the function for this endpoint
{import_statement}


def handler(job):
    """Handler for deployed QB endpoint. Accepts plain JSON kwargs."""
    job_input = job.get("input", {{}})
    try:
        result = {function_name}(**job_input)
        if inspect.iscoroutine(result):
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                loop = None
            if loop and loop.is_running():
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    result = pool.submit(asyncio.run, result).result()
            else:
                result = asyncio.run(result)
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {function_name}: %s",
            e,
            exc_info=True,
        )
        return {{"error": str(e), "traceback": traceback.format_exc()}}


if __name__ == "__main__":
    import runpod
    runpod.serverless.start({{"handler": handler}})
'''


class HandlerGenerator:
    """Generates handler_<name>.py files for each resource config."""

    def __init__(self, manifest: Union[Dict[str, Any], Manifest], build_dir: Path):
        self.manifest = manifest
        self.build_dir = build_dir

    def generate_handlers(self) -> List[Path]:
        """Generate all handler files for queue-based (non-LB) resources."""
        handler_paths = []

        # Handle both dict and Manifest types
        resources = (
            self.manifest.resources
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("resources", {})
        )

        for resource_name, resource_data in resources.items():
            # Skip load-balanced resources (handled by LBHandlerGenerator)
            # Use flag determined by isinstance() at scan time
            is_load_balanced = (
                resource_data.is_load_balanced
                if hasattr(resource_data, "is_load_balanced")
                else resource_data.get("is_load_balanced", False)
            )
            if is_load_balanced:
                continue

            handler_path = self._generate_handler(resource_name, resource_data)
            handler_paths.append(handler_path)

        return handler_paths

    def _generate_handler(self, resource_name: str, resource_data: Any) -> Path:
        """Generate a single handler file.

        Selects template based on resource type:
        - is_live_resource=True (flash run / local dev): HANDLER_TEMPLATE with
          cloudpickle-based create_handler() and full function registry
        - is_live_resource=False (deployed): DEPLOYED_HANDLER_TEMPLATE with
          plain JSON create_deployed_handler() and single function
        """
        handler_filename = f"handler_{resource_name}.py"
        handler_path = self.build_dir / handler_filename

        # Get timestamp from manifest
        timestamp = (
            self.manifest.generated_at
            if isinstance(self.manifest, Manifest)
            else self.manifest.get("generated_at", "")
        )

        # Get functions from resource (handle both dict and ResourceConfig)
        functions = (
            resource_data.functions
            if hasattr(resource_data, "functions")
            else resource_data.get("functions", [])
        )

        # Determine if this is a live resource (local dev) or deployed
        is_live_resource = (
            resource_data.is_live_resource
            if hasattr(resource_data, "is_live_resource")
            else resource_data.get("is_live_resource", False)
        )

        if is_live_resource:
            # Live resource (flash run): cloudpickle, multi-function registry
            imports = self._generate_imports(functions)
            registry = self._generate_registry(functions)
            handler_code = HANDLER_TEMPLATE.format(
                resource_name=resource_name,
                timestamp=timestamp,
                imports=imports,
                registry=registry,
            )
        else:
            # Deployed resource: plain JSON, single function per endpoint
            handler_code = self._generate_deployed_handler_code(
                resource_name, timestamp, functions
            )

        handler_path.write_text(handler_code)

        # Validate that generated handler can be imported
        self._validate_handler_imports(handler_path)

        return handler_path

    def _generate_deployed_handler_code(
        self,
        resource_name: str,
        timestamp: str,
        functions: List[Any],
    ) -> str:
        """Generate deployed handler code for a single-function endpoint.

        Args:
            resource_name: Name of the resource config.
            timestamp: Build timestamp.
            functions: List of function metadata for this resource.

        Returns:
            Generated handler source code.
        """
        if not functions:
            raise ValueError(
                f"Resource '{resource_name}' has no functions. "
                f"Cannot generate a deployed handler without at least one function."
            )

        # Use the first function (one function per deployed QB endpoint)
        func = functions[0]
        module = func.module if hasattr(func, "module") else func.get("module")
        name = func.name if hasattr(func, "name") else func.get("name")

        import_statement = (
            f"{name} = importlib.import_module('{module}').{name}"
            if module and name
            else "# No function to import"
        )

        return DEPLOYED_HANDLER_TEMPLATE.format(
            resource_name=resource_name,
            timestamp=timestamp,
            import_statement=import_statement,
            function_name=name or "None",
        )

    def _generate_imports(self, functions: List[Any]) -> str:
        """Generate import statements for functions using dynamic imports.

        Uses importlib.import_module() to handle module names with invalid
        Python identifiers (e.g., names starting with digits like '01_hello_world').
        """
        if not functions:
            return "# No functions to import"

        imports = []
        for func in functions:
            # Handle both dict and FunctionMetadata
            module = func.module if hasattr(func, "module") else func.get("module")
            name = func.name if hasattr(func, "name") else func.get("name")

            if module and name:
                # Use dynamic import to handle invalid identifiers
                imports.append(f"{name} = importlib.import_module('{module}').{name}")

        return "\n".join(imports) if imports else "# No functions to import"

    def _generate_registry(self, functions: List[Any]) -> str:
        """Generate function registry dictionary."""
        if not functions:
            return "    # No functions registered"

        registry_lines = []

        for func in functions:
            # Handle both dict and FunctionMetadata
            name = func.name if hasattr(func, "name") else func.get("name")
            registry_lines.append(f'    "{name}": {name},')

        return "\n".join(registry_lines)

    def _validate_handler_imports(self, handler_path: Path) -> None:
        """Validate that generated handler has valid Python syntax.

        Uses ast.parse() to check syntax without executing the module.
        This avoids ImportErrors from modules that only resolve at runtime
        inside Docker (e.g., numeric-prefixed module paths).

        Args:
            handler_path: Path to generated handler file

        Raises:
            ValueError: If handler has syntax errors
        """
        try:
            source = handler_path.read_text(encoding="utf-8")
            ast.parse(source, filename=str(handler_path))
        except SyntaxError as e:
            raise ValueError(f"Handler has syntax errors: {e}") from e
