################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
import asyncio
import logging
import uuid
from collections import deque
from concurrent.futures import Future
from typing import TYPE_CHECKING, Any, Callable, Dict, List

from typing_extensions import override

from flink_agents.api.agents.agent import Agent
from flink_agents.api.events.event import Event, InputEvent, OutputEvent
from flink_agents.api.events.tool_event import ToolRequestEvent
from flink_agents.api.memory.long_term_memory import BaseLongTermMemory
from flink_agents.api.memory_object import MemoryObject, MemoryType
from flink_agents.api.metric_group import MetricGroup
from flink_agents.api.resource import Resource, ResourceType
from flink_agents.api.runner_context import AsyncExecutionResult, RunnerContext
from flink_agents.plan.configuration import AgentConfiguration
from flink_agents.runtime.agent_runner import AgentRunner
from flink_agents.runtime.local_memory_object import LocalMemoryObject
from flink_agents.runtime.resource_cache import ResourceCache

if TYPE_CHECKING:
    from flink_agents.plan.agent_plan import AgentPlan

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class LocalRunnerContext(RunnerContext):
    """Implementation of RunnerContext for local agent execution.

    Attributes:
    ----------
    __agent_plan : AgentPlan
        Internal agent plan for this context.
    __key : Any
        Unique identifier for the context, correspond to the key in flink KeyedStream.
    events : deque[Event]
        Queue of events to be processed in this context.
    action_name: str
        Name of the action being executed.
    """

    __agent_plan: Any
    __key: Any
    events: deque[Event]
    action_name: str
    _sensory_mem_store: dict[str, Any]
    _short_term_mem_store: dict[str, Any]
    _sensory_memory: MemoryObject
    _short_term_memory: MemoryObject
    _config: AgentConfiguration

    def __init__(
        self, agent_plan: "AgentPlan", key: Any, config: AgentConfiguration
    ) -> None:
        """Initialize a new context with the given agent and key.

        Parameters
        ----------
        agent_plan : AgentPlan
            Agent plan used for this context.
        key : Any
            Unique context identifier, which is corresponding to the key in flink
            KeyedStream.
        """
        self.__agent_plan = agent_plan
        self.__resource_cache = ResourceCache(
            agent_plan.resource_providers, agent_plan.config
        )
        self.__key = key
        self.events = deque()
        self._sensory_mem_store = {}
        self._short_term_mem_store = {}
        self._sensory_memory = LocalMemoryObject(
            MemoryType.SENSORY, self._sensory_mem_store, LocalMemoryObject.ROOT_KEY
        )
        self._short_term_memory = LocalMemoryObject(
            MemoryType.SHORT_TERM,
            self._short_term_mem_store,
            LocalMemoryObject.ROOT_KEY,
        )
        self._config = config

    @property
    def key(self) -> Any:
        """Get the unique identifier for this context.

        Returns:
        -------
        Any
            The unique identifier for this context.
        """
        return self.__key

    @override
    def send_event(self, event: Event) -> None:
        """Send an event to the context's event queue and log it.

        Parameters
        ----------
        event : Event
            The event to be added to the queue.
        """
        logger.info("key: %s, send_event: %s", self.__key, event)
        self.events.append(event)

    @override
    def get_resource(
        self, name: str, type: ResourceType, metric_group: MetricGroup = None
    ) -> Resource:
        return self.__resource_cache.get_resource(name, type)

    @property
    @override
    def action_config(self) -> Dict[str, Any]:
        """Get config of the action."""
        return self.__agent_plan.get_action_config(action_name=self.action_name)

    @override
    def get_action_config_value(self, key: str) -> Any:
        """Get config option value of the key."""
        return self.__agent_plan.get_action_config_value(
            action_name=self.action_name, key=key
        )

    @property
    @override
    def sensory_memory(self) -> MemoryObject:
        """Get the sensory memory object associated with this context.

        Returns:
        -------
        MemoryObject
            The root object of the short-term memory.
        """
        return self._sensory_memory

    @property
    @override
    def short_term_memory(self) -> MemoryObject:
        """Get the short-term memory object associated with this context.

        Returns:
        -------
        MemoryObject
            The root object of the short-term memory.
        """
        return self._short_term_memory

    @property
    @override
    def long_term_memory(self) -> BaseLongTermMemory:
        err_msg = "Long-Term Memory is not supported for local agent execution yet."
        raise NotImplementedError(err_msg)

    @property
    @override
    def agent_metric_group(self) -> None:
        # TODO: Support metric mechanism for local agent execution.
        err_msg = "Metric mechanism is not supported for local agent execution yet."
        logger.warning(err_msg)
        return

    @property
    @override
    def action_metric_group(self) -> None:
        # TODO: Support metric mechanism for local agent execution.
        err_msg = "Metric mechanism is not supported for local agent execution yet."
        logger.warning(err_msg)
        return

    @override
    def durable_execute(
        self,
        func: Callable[[Any], Any],
        *args: Any,
        reconciler: Callable[[], Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        """Synchronously execute the provided function. Access to memory
        is prohibited within the function.

        Note: Local runner does not support durable execution, so recovery
        is not available.
        """
        logger.warning(
            "Local runner does not support durable execution; recovery is not available."
        )
        return func(*args, **kwargs)

    @override
    def durable_execute_async(
        self,
        func: Callable[[Any], Any],
        *args: Any,
        reconciler: Callable[[], Any] | None = None,
        **kwargs: Any,
    ) -> AsyncExecutionResult:
        """Asynchronously execute the provided function. Access to memory
        is prohibited within the function.

        Note: Local runner executes synchronously but returns an AsyncExecutionResult
        for API consistency. Durable execution is not supported.
        """
        logger.warning(
            "Local runner does not support durable execution; recovery is not available."
        )
        # Execute synchronously and wrap the result in a completed Future
        future: Future = Future()
        try:
            result = func(*args, **kwargs)
            future.set_result(result)
        except Exception as e:
            future.set_exception(e)

        # Create a mock executor that returns the pre-completed future
        class _SyncExecutor:
            def __init__(self, completed_future: Future) -> None:
                self._future = completed_future

            def submit(self, fn: Callable, *args: Any, **kwargs: Any) -> Future:
                return self._future

        return AsyncExecutionResult(_SyncExecutor(future), func, args, kwargs)

    @property
    @override
    def config(self) -> AgentConfiguration:
        return self._config

    def clear_sensory_memory(self) -> None:
        """Clean up sensory memory."""
        self._sensory_mem_store.clear()

    def close(self) -> None:
        """Cleanup the resource."""
        if self.__resource_cache is not None:
            try:
                self.__resource_cache.close()
            finally:
                self.__resource_cache = None


class LocalRunner(AgentRunner):
    """Agent runner implementation for local execution, which is
    convenient for debugging.

    Attributes:
    ----------
    __agent_plan : AgentPlan
        Internal agent plan.
    __keyed_contexts : dict[Any, LocalRunnerContext]
        Dictionary of active contexts indexed by key.
    __outputs:
        Outputs generated by agent execution.
    __tool_request_events:
        ToolRequestEvents observed during agent execution.
    __config:
        Internal configration.
    """

    __agent_plan: Any
    __keyed_contexts: Dict[Any, LocalRunnerContext]
    __outputs: List[Dict[str, Any]]
    __tool_request_events: List[ToolRequestEvent]
    __config: AgentConfiguration

    def __init__(self, agent: Agent, config: AgentConfiguration) -> None:
        """Initialize the runner with the provided agent.

        Parameters
        ----------
        agent : Agent
            The agent class to convert and run.
        """
        from flink_agents.plan.agent_plan import AgentPlan

        self.__agent_plan = AgentPlan.from_agent(agent, config)
        self.__keyed_contexts = {}
        self.__outputs = []
        self.__tool_request_events = []
        self.__config = config

    @override
    def run(self, **data: Dict[str, Any]) -> Any:
        """Execute the agent to process the given data.

        Parameters
        ----------
        **data : dict[str, Any]
            input record from upstream.

        Returns:
        -------
        key
            The key of the input that was processed.
        """
        if "key" in data:
            key = data["key"]
        elif "k" in data:
            key = data["k"]
        else:
            key = uuid.uuid4()

        if key not in self.__keyed_contexts:
            self.__keyed_contexts[key] = LocalRunnerContext(
                self.__agent_plan, key, self.__config
            )
        context = self.__keyed_contexts[key]
        context.clear_sensory_memory()

        if "value" in data:
            input_event = InputEvent(input=data["value"])
        elif "v" in data:
            input_event = InputEvent(input=data["v"])
        else:
            msg = "Input data must be dict has 'v' or 'value' field"
            raise RuntimeError(msg)

        context.send_event(input_event)

        while len(context.events) > 0:
            event = context.events.popleft()
            if isinstance(event, OutputEvent):
                self.__outputs.append({key: event.output})
                continue
            if isinstance(event, ToolRequestEvent):
                self.__tool_request_events.append(event)
                # Fall through: the request must still dispatch to its action.
            event_type = event.get_type()
            for action in self.__agent_plan.get_actions(event_type):
                logger.info("key: %s, performing action: %s", key, action.name)
                context.action_name = action.name
                func_result = action.exec(event, context)
                if asyncio.iscoroutine(func_result):
                    # Drive the coroutine to completion using send()
                    try:
                        while True:
                            func_result.send(None)
                    except StopIteration:
                        # Coroutine completed normally
                        pass
                    except Exception:
                        logger.exception("Error in async execution")
                        raise
        return key

    def get_outputs(self) -> List[Dict[str, Any]]:
        """Get the outputs generated by agent execution.

        Returns:
        -------
        List[Dict[str, Any]]
            The agent execution outputs.
        """
        return self.__outputs

    def get_tool_request_events(self) -> List[ToolRequestEvent]:
        """Get the ToolRequestEvents captured during agent execution.

        Returns:
        -------
        List[ToolRequestEvent]
            The ToolRequestEvents observed in the run loop, in order.
        """
        return self.__tool_request_events
