"""UiPath Function Runtime - wrapper for executing Python functions."""

import importlib.util
import inspect
import logging
import sys
import uuid
from dataclasses import is_dataclass
from pathlib import Path
from types import ModuleType
from typing import Any, AsyncGenerator, Callable, Type, cast, get_type_hints

from uipath.runtime import (
    UiPathExecuteOptions,
    UiPathRuntimeEvent,
    UiPathRuntimeResult,
    UiPathRuntimeStatus,
)
from uipath.runtime.errors import (
    UiPathErrorCategory,
    UiPathErrorCode,
    UiPathRuntimeError,
)
from uipath.runtime.events import UiPathRuntimeStateEvent, UiPathRuntimeStatePhase
from uipath.runtime.schema import UiPathRuntimeSchema, transform_attachments

from .graph_builder import build_call_graph
from .schema_gen import get_type_schema
from .type_conversion import (
    convert_from_class,
    convert_to_class,
    is_pydantic_model,
)

logger = logging.getLogger(__name__)


class UiPathFunctionsRuntime:
    """Runtime wrapper for a single Python function with full script executor compatibility."""

    def __init__(
        self,
        file_path: str,
        function_name: str,
        entrypoint_name: str,
        entrypoint_type: str = "function",
    ):
        """Initialize the function runtime.

        Args:
            file_path: Path to the Python file containing the function
            function_name: Name of the function to execute
            entrypoint_name: Name of the entrypoint
            entrypoint_type: Type of entrypoint - 'function' or 'agent'
        """
        self.file_path = Path(file_path)
        self.function_name = function_name
        self.entrypoint_name = entrypoint_name
        self.entrypoint_type = entrypoint_type
        self._function: Callable[..., Any] | None = None
        self._module: ModuleType | None = None

    def _load_module(self) -> None:
        """Load the Python module containing the function."""
        if self._module is not None:
            return

        spec = importlib.util.spec_from_file_location(
            "dynamic_module", str(self.file_path)
        )
        if not spec or not spec.loader:
            raise UiPathRuntimeError(
                UiPathErrorCode.IMPORT_ERROR,
                "Module import failed",
                f"Could not load spec for {self.file_path}",
                UiPathErrorCategory.USER,
            )

        module = importlib.util.module_from_spec(spec)
        sys.modules["dynamic_module"] = module
        self._module = module

        try:
            spec.loader.exec_module(module)
        except Exception as e:
            raise UiPathRuntimeError(
                UiPathErrorCode.MODULE_EXECUTION_ERROR,
                "Module execution failed",
                f"Error executing module: {e}",
                UiPathErrorCategory.USER,
            ) from e

    def _load_function(self) -> Callable[..., Any]:
        """Load the function from the module."""
        if self._function is not None:
            return self._function

        self._load_module()

        if not hasattr(self._module, self.function_name):
            raise UiPathRuntimeError(
                UiPathErrorCode.ENTRYPOINT_FUNCTION_MISSING,
                "Function not found",
                f"Function '{self.function_name}' not found in {self.file_path}",
                UiPathErrorCategory.USER,
            )

        self._function = getattr(self._module, self.function_name)
        return self._function

    async def _execute_function(
        self, func: Callable[..., Any], input_data: dict[str, Any]
    ) -> dict[str, Any]:
        """Execute function with proper input conversion and error handling."""
        # Unwrap to inspect the original signature for type-based conversion,
        # but still call the outer (decorated) func for proper tracing/hooks.
        unwrapped = inspect.unwrap(func)
        sig = inspect.signature(unwrapped)
        params = list(sig.parameters.values())
        is_async = inspect.iscoroutinefunction(unwrapped)

        # No parameters - call without args
        if not params:
            result = await func() if is_async else func()
            return convert_from_class(result) if result is not None else {}

        # Get first parameter info.
        # Use get_type_hints() to resolve PEP 563 stringified annotations
        # (from `from __future__ import annotations`).  inspect.signature()
        # returns raw strings in that case, which breaks type detection.
        input_param = params[0]
        try:
            hints = get_type_hints(unwrapped)
            input_type = hints.get(input_param.name, inspect.Parameter.empty)
        except Exception:
            input_type = input_param.annotation

        # Typed parameter (class, dataclass, or Pydantic)
        if input_type != inspect.Parameter.empty and (
            is_dataclass(input_type)
            or is_pydantic_model(input_type)
            or (inspect.isclass(input_type) and hasattr(input_type, "__annotations__"))
        ):
            typed_input = convert_to_class(input_data, cast(Type[Any], input_type))
            result = await func(typed_input) if is_async else func(typed_input)
        else:
            # Dict/untyped parameter
            result = await func(input_data) if is_async else func(input_data)

        return convert_from_class(result) if result is not None else {}

    async def execute(
        self,
        input: dict[str, Any] | None = None,
        options: UiPathExecuteOptions | None = None,
    ) -> UiPathRuntimeResult:
        """Execute the function with the given input."""
        try:
            func = self._load_function()
            output = await self._execute_function(func, input or {})

            return UiPathRuntimeResult(
                output=output,
                status=UiPathRuntimeStatus.SUCCESSFUL,
            )

        except UiPathRuntimeError:
            raise
        except Exception as e:
            raise UiPathRuntimeError(
                UiPathErrorCode.FUNCTION_EXECUTION_ERROR,
                f"Function execution failed: {self.function_name}",
                str(e),
                UiPathErrorCategory.USER,
            ) from e

    async def stream(
        self,
        input: dict[str, Any] | None = None,
        options: UiPathExecuteOptions | None = None,
    ) -> AsyncGenerator[UiPathRuntimeEvent, None]:
        """Stream execution results with lifecycle phase events.

        Yields STARTED → COMPLETED/FAULTED → result so consumers can
        observe the execution lifecycle.
        """
        yield UiPathRuntimeStateEvent(
            node_name=self.function_name,
            phase=UiPathRuntimeStatePhase.STARTED,
            payload=input or {},
        )

        try:
            result = await self.execute(input, options)
        except UiPathRuntimeError as e:
            yield UiPathRuntimeStateEvent(
                node_name=self.function_name,
                phase=UiPathRuntimeStatePhase.FAULTED,
                payload={"error": str(e)},
            )
            raise

        output = result.output
        if output is None:
            completed_payload: dict[str, Any] = {}
        elif isinstance(output, dict):
            completed_payload = output
        else:
            completed_payload = {"output": str(output)}
        yield UiPathRuntimeStateEvent(
            node_name=self.function_name,
            phase=UiPathRuntimeStatePhase.COMPLETED,
            payload=completed_payload,
        )

        yield result

    async def get_schema(self) -> UiPathRuntimeSchema:
        """Get schema for the function."""
        func = self._load_function()

        # Unwrap decorated functions to get the original signature and type hints.
        # This handles decorators (e.g. @traced) that use functools.wraps.
        unwrapped = inspect.unwrap(func)
        hints = get_type_hints(unwrapped)
        sig = inspect.signature(unwrapped)

        # Determine input schema
        if not sig.parameters:
            input_schema = {}
        else:
            input_param_name = next(iter(sig.parameters))
            raw_input_schema = get_type_schema(hints.get(input_param_name))
            input_schema = transform_attachments(raw_input_schema)

        # Determine output schema
        raw_output_schema = get_type_schema(hints.get("return"))
        output_schema = transform_attachments(raw_output_schema)

        # Build call graph from AST
        graph = None
        try:
            graph = build_call_graph(
                str(self.file_path),
                self.function_name,
                project_dir=str(self.file_path.parent),
            )
        except Exception:
            logger.debug(
                "Failed to build call graph for %s:%s",
                self.file_path,
                self.function_name,
                exc_info=True,
            )

        return UiPathRuntimeSchema(
            filePath=self.entrypoint_name,
            uniqueId=str(uuid.uuid4()),
            type=self.entrypoint_type,
            input=input_schema,
            output=output_schema,
            graph=graph,
        )

    async def dispose(self) -> None:
        """Cleanup resources."""
        self._function = None
        if "dynamic_module" in sys.modules:
            del sys.modules["dynamic_module"]
        self._module = None
