"""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__)

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}})
'''

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

Deployed endpoint handler for a @remote class. Instantiates the class
once at cold start, then dispatches to methods per request.

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 class for this endpoint
{import_statement}

# instantiate once at cold start
_instance = {class_name}()

# public methods available for dispatch
_METHODS = {methods_dict}


def _run_maybe_async(result):
    """resolve a possibly-async result to a plain value."""
    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:
                return pool.submit(asyncio.run, result).result()
        return asyncio.run(result)
    return result


def handler(job):
    """Handler for deployed class-based QB endpoint.

    Dispatches to a method on the singleton class instance.
    If the class has exactly one public method, input is passed
    directly as kwargs. If multiple methods exist, input must
    include a "method" key to select the target.
    """
    job_input = job.get("input", {{}})
    try:
        if len(_METHODS) == 1:
            method_name = next(iter(_METHODS))
        else:
            method_name = job_input.pop("method", None)
            if method_name is None:
                return {{
                    "error": (
                        "class {class_name} has multiple methods: "
                        + ", ".join(sorted(_METHODS))
                        + ". include a \\"method\\" key in the input."
                    )
                }}
            if method_name not in _METHODS:
                return {{
                    "error": (
                        f"unknown method '{{method_name}}' on {class_name}. "
                        f"available: {{', '.join(sorted(_METHODS))}}"
                    )
                }}

        method = getattr(_instance, method_name)
        result = _run_maybe_async(method(**job_input))
        return result
    except Exception as e:
        _logger.error(
            "Deployed handler error for {class_name}.%s: %s",
            job_input.get("method", "<default>"),
            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 handler file for a deployed QB endpoint.

        Produces a plain-JSON handler that accepts raw dicts (no cloudpickle).
        For classes, instantiates once at cold start and dispatches to methods.
        For functions, calls the function directly with job input as kwargs.
        """
        handler_filename = f"handler_{resource_name}.py"
        handler_path = self.build_dir / handler_filename

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

        functions = (
            resource_data.functions
            if hasattr(resource_data, "functions")
            else resource_data.get("functions", [])
        )

        handler_code = self._generate_deployed_handler_code(
            resource_name, timestamp, functions
        )

        handler_path.write_text(handler_code)
        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 QB endpoint.

        Selects the class template for class-based workers and the
        function template for plain functions.
        """
        if not functions:
            raise ValueError(
                f"Resource '{resource_name}' has no functions. "
                f"Cannot generate a deployed handler without at least one function."
            )

        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")
        is_class = (
            func.is_class if hasattr(func, "is_class") else func.get("is_class", False)
        )

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

        if is_class:
            class_methods = (
                func.class_methods
                if hasattr(func, "class_methods")
                else func.get("class_methods", [])
            )
            methods_dict = {m: m for m in class_methods} if class_methods else {}
            return DEPLOYED_CLASS_HANDLER_TEMPLATE.format(
                resource_name=resource_name,
                timestamp=timestamp,
                import_statement=import_statement,
                class_name=name or "None",
                methods_dict=repr(methods_dict),
            )

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

    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
