################################################################################
#  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 importlib
from abc import ABC, abstractmethod
from enum import Enum
from typing import TYPE_CHECKING, Any, Dict, Type

from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator

from flink_agents.api.resource_context import ResourceContext

if TYPE_CHECKING:
    from flink_agents.api.metric_group import MetricGroup


class ResourceType(Enum):
    """Type enum of resource.

    Currently, support chat_model, chat_model_server, tool, embedding_model,
    vector_store, prompt, mcp_server, skills.
    """

    CHAT_MODEL = "chat_model"
    CHAT_MODEL_CONNECTION = "chat_model_connection"
    TOOL = "tool"
    EMBEDDING_MODEL = "embedding_model"
    EMBEDDING_MODEL_CONNECTION = "embedding_model_connection"
    VECTOR_STORE = "vector_store"
    PROMPT = "prompt"
    MCP_SERVER = "mcp_server"
    SKILLS = "skills"


class Resource(BaseModel, ABC):
    """Base abstract class of all kinds of resources, includes chat model,
    prompt, tools and so on.

    Resource extends BaseModel only for decreasing the complexity of attribute
    declaration of subclasses, this not represents Resource object is serializable.

    Attributes:
    ----------
    resource_context : ResourceContext
        Get other resource object declared in the same Agent. The first argument is
        resource name and the second argument is resource type.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    resource_context: ResourceContext | None = Field(exclude=True, default=None)

    # The metric group bound to this resource, injected in RunnerContext#get_resource
    _metric_group: "MetricGroup | None" = PrivateAttr(default=None)

    @classmethod
    @abstractmethod
    def resource_type(cls) -> ResourceType:
        """Return resource type of class."""

    def set_metric_group(self, metric_group: "MetricGroup") -> None:
        """Set the metric group for this resource.

        Parameters
        ----------
        metric_group : MetricGroup
            The metric group to bind.
        """
        self._metric_group = metric_group

    @property
    def metric_group(self) -> "MetricGroup | None":
        """Get the bound metric group.

        Returns:
        -------
        MetricGroup | None
            The bound metric group, or None if not set.
        """
        return self._metric_group

    def open(self) -> None:
        """Open the resource."""

    def close(self) -> None:
        """Close the resource."""


class SerializableResource(Resource, ABC):
    """Resource which is serializable."""

    @model_validator(mode="after")
    def validate_serializable(self) -> "SerializableResource":
        """Ensure resource is serializable."""
        self.model_dump_json()
        return self


class ResourceDescriptor(BaseModel):
    """Descriptor for Resource instances, storing metadata for serialization and
    instantiation.

    Attributes:
        target_module: The module name of the resource class.
        target_clazz: The class name of the resource.
        arguments: Dictionary containing resource initialization parameters.
    """

    _clazz: Type[Resource] = None
    target_module: str
    target_clazz: str
    arguments: Dict[str, Any]

    def __init__(
        self,
        /,
        *,
        clazz: str | None = None,
        target_module: str | None = None,
        target_clazz: str | None = None,
        arguments: Dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        """Initialize ResourceDescriptor.

        Args:
            clazz: The fully qualified name of the resource implementation, including
                    module and class.
            target_module: The module name of the resource class.
            target_clazz: The class name of the resource.
            arguments: Dictionary containing resource initialization parameters.
            **kwargs: Additional keywords arguments for resource initialization,
            will merge into arguments.

        Usage:
            descriptor = ResourceDescriptor(clazz="flink_agents.integrations.chat_models
                                            .ollama_chat_model.OllamaChatModelConnection",
                                            param1="value1",
                                            param2="value2")
        """
        if clazz is not None:
            parts = clazz.split(".")
            target_module = ".".join(parts[:-1])
            target_clazz = parts[-1]

        if target_clazz is None or target_module is None:
            msg = "The fully qualified name of the resource must be specified"
            raise ValueError(msg)

        args = {}
        if arguments is not None:
            args.update(arguments)
        args.update(kwargs)

        super().__init__(
            target_module=target_module, target_clazz=target_clazz, arguments=args
        )

    @property
    def clazz(self) -> Type[Resource]:
        """Get the class of the resource."""
        if self._clazz is None:
            module = importlib.import_module(self.target_module)
            self._clazz = getattr(module, self.target_clazz)
        return self._clazz

    def __eq__(self, other: object) -> bool:
        """Compare ResourceDescriptor objects, ignoring private _clazz field.

        This ensures that deserialized objects (with _clazz=None) can be compared
        equal to runtime objects (with _clazz set) as long as their serializable
        fields match.
        """
        if not isinstance(other, ResourceDescriptor):
            return False
        return (
            self.target_module == other.target_module
            and self.target_clazz == other.target_clazz
            and self.arguments == other.arguments
        )

    def __hash__(self) -> int:
        """Generate hash for ResourceDescriptor."""
        return hash(
            (
                self.target_module,
                self.target_clazz,
                tuple(sorted(self.arguments.items())),
            )
        )


def get_resource_class(module_path: str, class_name: str) -> Type[Resource]:
    """Get Resource class from separate module path and class name.

    Args:
        module_path: Python module path (e.g., 'your.module.path').
        class_name: Class name (e.g., 'YourResourceClass').

    Returns:
        The Resource class type.
    """
    module = importlib.import_module(module_path)
    return getattr(module, class_name)


class ResourceName:
    """Hierarchical resource class names for pointing a resource implementation in
    ResourceDescriptor.

    Structure:
        - Python implementation: ResourceType.PROVIDER_RESOURCEKIND
        - Java implementation: ResourceType.Java.PROVIDER_RESOURCEKIND

    Example usage:
        # Python implementation
        ResourceName.ChatModel.OLLAMA_CONNECTION
        ResourceName.ChatModel.OPENAI_COMPLETIONS_SETUP

        # Java implementation
        ResourceName.ChatModel.Java.OLLAMA_CONNECTION
    """

    class ChatModel:
        """ChatModel resource names."""

        # Anthropic
        ANTHROPIC_CONNECTION = "flink_agents.integrations.chat_models.anthropic.anthropic_chat_model.AnthropicChatModelConnection"
        ANTHROPIC_SETUP = "flink_agents.integrations.chat_models.anthropic.anthropic_chat_model.AnthropicChatModelSetup"

        # Azure OpenAI
        AZURE_OPENAI_CONNECTION = "flink_agents.integrations.chat_models.azure.azure_openai_chat_model.AzureOpenAIChatModelConnection"
        AZURE_OPENAI_SETUP = "flink_agents.integrations.chat_models.azure.azure_openai_chat_model.AzureOpenAIChatModelSetup"

        # Ollama
        OLLAMA_CONNECTION = "flink_agents.integrations.chat_models.ollama_chat_model.OllamaChatModelConnection"
        OLLAMA_SETUP = "flink_agents.integrations.chat_models.ollama_chat_model.OllamaChatModelSetup"

        # OpenAI
        OPENAI_COMPLETIONS_CONNECTION = "flink_agents.integrations.chat_models.openai.openai_chat_model.OpenAIChatModelConnection"
        OPENAI_COMPLETIONS_SETUP = "flink_agents.integrations.chat_models.openai.openai_chat_model.OpenAIChatModelSetup"

        # Tongyi
        TONGYI_CONNECTION = "flink_agents.integrations.chat_models.tongyi_chat_model.TongyiChatModelConnection"
        TONGYI_SETUP = "flink_agents.integrations.chat_models.tongyi_chat_model.TongyiChatModelSetup"

        # Java Wrapper
        JAVA_WRAPPER_CONNECTION = (
            "flink_agents.api.chat_models.java_chat_model.JavaChatModelConnection"
        )
        JAVA_WRAPPER_SETUP = (
            "flink_agents.api.chat_models.java_chat_model.JavaChatModelSetup"
        )

        class Java:
            """Java implementations of ChatModel."""

            # Anthropic
            ANTHROPIC_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.anthropic.AnthropicChatModelConnection"
            ANTHROPIC_SETUP = "org.apache.flink.agents.integrations.chatmodels.anthropic.AnthropicChatModelSetup"

            # Azure
            AZURE_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.azureai.AzureAIChatModelConnection"
            AZURE_SETUP = "org.apache.flink.agents.integrations.chatmodels.azureai.AzureAIChatModelSetup"

            # Bedrock
            BEDROCK_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.bedrock.BedrockChatModelConnection"
            BEDROCK_SETUP = "org.apache.flink.agents.integrations.chatmodels.bedrock.BedrockChatModelSetup"

            # Ollama
            OLLAMA_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.ollama.OllamaChatModelConnection"
            OLLAMA_SETUP = "org.apache.flink.agents.integrations.chatmodels.ollama.OllamaChatModelSetup"

            # OpenAI Completions
            OPENAI_COMPLETIONS_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.openai.OpenAICompletionsConnection"
            OPENAI_COMPLETIONS_SETUP = "org.apache.flink.agents.integrations.chatmodels.openai.OpenAICompletionsSetup"

            OPENAI_RESPONSES_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.openai.OpenAIResponsesModelConnection"
            OPENAI_RESPONSES_SETUP = "org.apache.flink.agents.integrations.chatmodels.openai.OpenAIResponsesModelSetup"

            # Azure OpenAI
            AZURE_OPENAI_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.openai.AzureOpenAIChatModelConnection"
            AZURE_OPENAI_SETUP = "org.apache.flink.agents.integrations.chatmodels.openai.AzureOpenAIChatModelSetup"

    class EmbeddingModel:
        """EmbeddingModel resource names."""

        # Ollama
        OLLAMA_CONNECTION = "flink_agents.integrations.embedding_models.local.ollama_embedding_model.OllamaEmbeddingModelConnection"
        OLLAMA_SETUP = "flink_agents.integrations.embedding_models.local.ollama_embedding_model.OllamaEmbeddingModelSetup"

        # OpenAI
        OPENAI_CONNECTION = "flink_agents.integrations.embedding_models.openai_embedding_model.OpenAIEmbeddingModelConnection"
        OPENAI_SETUP = "flink_agents.integrations.embedding_models.openai_embedding_model.OpenAIEmbeddingModelSetup"

        # Tongyi
        TONGYI_CONNECTION = "flink_agents.integrations.embedding_models.tongyi_embedding_model.TongyiEmbeddingModelConnection"
        TONGYI_SETUP = "flink_agents.integrations.embedding_models.tongyi_embedding_model.TongyiEmbeddingModelSetup"

        # Java Wrapper
        JAVA_WRAPPER_CONNECTION = "flink_agents.api.embedding_models.java_embedding_model.JavaEmbeddingModelConnection"
        JAVA_WRAPPER_SETUP = "flink_agents.api.embedding_models.java_embedding_model.JavaEmbeddingModelSetup"

        class Java:
            """Java implementations of EmbeddingModel."""

            # Ollama
            OLLAMA_CONNECTION = "org.apache.flink.agents.integrations.embeddingmodels.ollama.OllamaEmbeddingModelConnection"
            OLLAMA_SETUP = "org.apache.flink.agents.integrations.embeddingmodels.ollama.OllamaEmbeddingModelSetup"

            # Bedrock
            BEDROCK_CONNECTION = "org.apache.flink.agents.integrations.embeddingmodels.bedrock.BedrockEmbeddingModelConnection"
            BEDROCK_SETUP = "org.apache.flink.agents.integrations.embeddingmodels.bedrock.BedrockEmbeddingModelSetup"

    class VectorStore:
        """VectorStore resource names."""

        # Chroma
        CHROMA_VECTOR_STORE = "flink_agents.integrations.vector_stores.chroma.chroma_vector_store.ChromaVectorStore"

        # Mem0 (gateway to Mem0's native vector stores: pgvector, milvus, qdrant,
        # redis, ...)
        MEM0_VECTOR_STORE = "flink_agents.integrations.vector_stores.mem0.mem0_vector_store.Mem0VectorStore"

        # Java Wrapper
        JAVA_WRAPPER_VECTOR_STORE = (
            "flink_agents.api.vector_stores.java_vector_store.JavaVectorStore"
        )
        JAVA_WRAPPER_COLLECTION_MANAGEABLE_VECTOR_STORE = "flink_agents.api.vector_stores.java_vector_store.JavaCollectionManageableVectorStore"

        class Java:
            """Java implementations of VectorStore."""

            # Elasticsearch
            ELASTICSEARCH_VECTOR_STORE = "org.apache.flink.agents.integrations.vectorstores.elasticsearch.ElasticsearchVectorStore"

            # Amazon OpenSearch (Serverless or Service domains)
            OPENSEARCH_VECTOR_STORE = "org.apache.flink.agents.integrations.vectorstores.opensearch.OpenSearchVectorStore"

            # Amazon S3 Vectors
            S3_VECTORS_VECTOR_STORE = "org.apache.flink.agents.integrations.vectorstores.s3vectors.S3VectorsVectorStore"

            # Milvus
            MILVUS_VECTOR_STORE = "org.apache.flink.agents.integrations.vectorstores.milvus.MilvusVectorStore"

    # MCP resource names
    MCP_SERVER = "flink_agents.integrations.mcp.mcp.MCPServer"
