# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import inspect
import threading
import weakref
from typing import (
    TYPE_CHECKING,
    Any,
    Awaitable,
    ClassVar,
    Dict,
    Optional,
    Union,
)

from mantisdk.adapter import TraceAdapter
from mantisdk.client import MantisdkClient
from mantisdk.store.base import LightningStore
from mantisdk.types import Dataset, NamedResources

if TYPE_CHECKING:
    from mantisdk.llm_proxy import LLMProxy
    from mantisdk.trainer import Trainer


class Algorithm:
    """Algorithm is the strategy, or tuner to train the agent.

    Subclasses should set the class-level metadata attributes so the platform
    can auto-discover algorithms and sync them to the database.
    """

    # Platform metadata — synced to DB via auto-discovery (__subclasses__)
    name: ClassVar[str] = ""
    display_name: ClassVar[str] = ""
    description: ClassVar[str] = ""
    requirements: ClassVar[Dict[str, Any]] = {}
    config_schema: ClassVar[Optional[Dict[str, Any]]] = None
    manages_containers: ClassVar[bool] = False
    trainer_resources: ClassVar[Dict[str, Any]] = {}
    """Container resources for the algorithm trainer process (reserved for Phase 4).

    Example: ``{"gpu": 1, "memory": "8g", "cpu": 2.0}``

    Used by the orchestrator to size the trainer container when
    ``manages_containers=False``.  Currently unused -- reserved for verl
    multi-container support.
    """

    services: ClassVar[Dict[str, Any]] = {}
    """Sidecar services the algorithm requires (reserved for Phase 4).

    Example: ``{"vllm": {"image": "vllm/vllm:latest", "gpu": 1}}``

    Used by Layer 2 orchestration for algorithms that need inference servers
    alongside the trainer.  Currently unused -- reserved for verl.
    """

    # Runtime state (set by Trainer)
    _trainer_ref: weakref.ReferenceType[Trainer] | None = None
    _llm_proxy_ref: weakref.ReferenceType["LLMProxy"] | None = None
    _store: LightningStore | None = None
    _initial_resources: NamedResources | None = None
    _adapter_ref: weakref.ReferenceType[TraceAdapter[Any]] | None = None

    # Cancellation support -- thread-safe event so the platform worker can
    # signal synchronous algorithms running in a background thread.
    _cancel_event: threading.Event

    def __init__(self) -> None:
        self._cancel_event = threading.Event()

    def request_cancel(self) -> None:
        """Signal the algorithm to stop as soon as possible.

        Thread-safe: can be called from the async event loop while the
        algorithm runs in a worker thread.  Subclasses that manage their
        own engine (e.g. GEPA) should override this to also stop the
        inner loop.
        """
        self._cancel_event.set()

    @property
    def is_cancelled(self) -> bool:
        """Check whether cancellation has been requested."""
        return self._cancel_event.is_set()

    def is_async(self) -> bool:
        """Return True if the algorithm is asynchronous."""
        return inspect.iscoroutinefunction(self.run)

    def set_trainer(self, trainer: Trainer) -> None:
        """
        Set the trainer for this algorithm.

        Args:
            trainer: The Trainer instance that will handle training and validation.
        """
        self._trainer_ref = weakref.ref(trainer)

    def get_trainer(self) -> Trainer:
        """
        Get the trainer for this algorithm.

        Returns:
            The Trainer instance associated with this agent.
        """
        if self._trainer_ref is None:
            raise ValueError("Trainer has not been set for this agent.")
        trainer = self._trainer_ref()
        if trainer is None:
            raise ValueError("Trainer reference is no longer valid (object has been garbage collected).")
        return trainer

    def set_llm_proxy(self, llm_proxy: LLMProxy | None) -> None:
        """
        Set the LLM proxy for this algorithm to reuse when available.

        Args:
            llm_proxy: The LLMProxy instance configured by the trainer, if any.
        """
        self._llm_proxy_ref = weakref.ref(llm_proxy) if llm_proxy is not None else None

    def get_llm_proxy(self) -> Optional[LLMProxy]:
        """
        Retrieve the configured LLM proxy instance, if one has been set.

        Returns:
            The active LLMProxy instance or None when not configured.
        """
        if self._llm_proxy_ref is None:
            return None

        llm_proxy = self._llm_proxy_ref()
        if llm_proxy is None:
            raise ValueError("LLM proxy reference is no longer valid (object has been garbage collected).")

        return llm_proxy

    def set_adapter(self, adapter: TraceAdapter[Any]) -> None:
        """
        Set the adapter for this algorithm to collect and convert traces.
        """
        self._adapter_ref = weakref.ref(adapter)

    def get_adapter(self) -> TraceAdapter[Any]:
        """
        Retrieve the adapter for this algorithm to communicate with the runners.
        """
        if self._adapter_ref is None:
            raise ValueError("Adapter has not been set for this algorithm.")
        adapter = self._adapter_ref()
        if adapter is None:
            raise ValueError("Adapter reference is no longer valid (object has been garbage collected).")
        return adapter

    def set_store(self, store: LightningStore) -> None:
        """
        Set the store for this algorithm to communicate with the runners.

        Store is set directly instead of using weakref because its copy is meant to be
        maintained throughout the algorithm's lifecycle.
        """
        self._store = store

    def get_store(self) -> LightningStore:
        """
        Retrieve the store for this algorithm to communicate with the runners.
        """
        if self._store is None:
            raise ValueError("Store has not been set for this algorithm.")
        return self._store

    def get_initial_resources(self) -> Optional[NamedResources]:
        """
        Get the initial resources for this algorithm.
        """
        return self._initial_resources

    def set_initial_resources(self, resources: NamedResources) -> None:
        """
        Set the initial resources for this algorithm.
        """
        self._initial_resources = resources

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return self.run(*args, **kwargs)

    def run(
        self,
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> Union[None, Awaitable[None]]:
        """Subclasses should implement this method to implement the algorithm.

        Args:
            train_dataset: The dataset to train on. Not all algorithms require a training dataset.
            val_dataset: The dataset to validate on. Not all algorithms require a validation dataset.

        Returns:
            Algorithm should refrain from returning anything. It should just run the algorithm.
        """
        raise NotImplementedError("Subclasses must implement run().")

    def verify(self, config: Dict[str, Any]) -> Dict[str, Any]:
        """Validate algorithm configuration before a run starts.

        When ``config_schema`` is set (JSON Schema dict), the config is
        validated against it automatically.  Subclasses may override for
        additional algorithm-specific checks.

        Args:
            config: The algorithm configuration to validate.

        Returns:
            Dict with at least ``{"valid": bool}``.  May include
            ``"smart_defaults"`` (defaults derived from schema) and
            ``"errors"`` (list of validation error strings).
        """
        if self.config_schema is not None:
            try:
                import jsonschema

                jsonschema.validate(instance=config, schema=self.config_schema)
            except ImportError:
                pass  # jsonschema not installed — skip validation
            except jsonschema.ValidationError as exc:
                return {
                    "valid": False,
                    "errors": [exc.message],
                    "path": list(exc.absolute_path),
                }

        # Derive smart defaults from schema property defaults
        smart_defaults: Dict[str, Any] = {}
        if self.config_schema is not None:
            for prop, prop_schema in self.config_schema.get("properties", {}).items():
                if prop not in config and "default" in prop_schema:
                    smart_defaults[prop] = prop_schema["default"]

        return {"valid": True, "smart_defaults": smart_defaults}

    def get_client(self) -> Optional[MantisdkClient]:
        """Get the client to communicate with the algorithm.

        .. deprecated::
            ``get_client()`` is deprecated and will be removed in a future
            version.  Use :meth:`get_store` for communication with runners.

        Returns:
            None.  Subclasses that still rely on MantisdkClient may override.
        """
        import warnings

        warnings.warn(
            "Algorithm.get_client() is deprecated and will be removed. " "Use Algorithm.get_store() instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        return None
