#
# Copyright 2026 The Dapr Authors
# Licensed 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.
#

from __future__ import annotations

import logging
import re
from os import getenv
from enum import StrEnum
from dataclasses import dataclass, field
from typing import (
    Any,
    Callable,
    Dict,
    List,
    MutableMapping,
    Optional,
    Sequence,
    Type,
    Union,
)

from pydantic import BaseModel, Field

from dapr_agents.agents.schemas import (
    AgentWorkflowEntry,
    AgentWorkflowMessage,
)

from dapr_agents.memory import ConversationListMemory, MemoryBase
from dapr_agents.storage.daprstores.stateservice import StateStoreService

_JINJA_PLACEHOLDER_PATTERN = re.compile(r"(?<!\{)\{\s*(\w+)\s*\}(?!\})")

# JSON Schema export constants
_JSON_SCHEMA_KEY = "$schema"
_JSON_SCHEMA_DRAFT_URL = "https://json-schema.org/draft/2020-12/schema"
_JSON_SCHEMA_VERSION_KEY = "version"


def _ensure_jinja_placeholders(text: str) -> str:
    return _JINJA_PLACEHOLDER_PATTERN.sub(r"{{\1}}", text)


def _empty_headers() -> Dict[str, str]:
    return {}


# Type hooks for state customization
EntryFactory = Callable[..., Any]
MessageCoercer = Callable[[Dict[str, Any]], Any]
EntryContainerGetter = Callable[[BaseModel], Optional[MutableMapping[str, Any]]]


@dataclass
class StateModelBundle:
    """
    Bundled state schema configuration for an agent/orchestrator type.

    With one-key-per-workflow, each state store key holds a single workflow
    entry (entry_model_cls). This bundle identifies that type and related hooks.

    Attributes:
        entry_model_cls: Pydantic model class for one workflow's state (per key).
        message_model_cls: Pydantic model class for workflow/system messages.
        entry_factory: Optional factory to create workflow entry instances.
        message_coercer: Optional function to transform message dicts.
    """

    entry_model_cls: Type[BaseModel]
    message_model_cls: Type[BaseModel]
    entry_factory: Optional[EntryFactory] = None
    message_coercer: Optional[MessageCoercer] = None


DEFAULT_AGENT_WORKFLOW_BUNDLE = StateModelBundle(
    entry_model_cls=AgentWorkflowEntry,
    message_model_cls=AgentWorkflowMessage,
)


@dataclass
class WorkflowGrpcOptions:
    """
    Optional overrides for Durable Task gRPC channel limits.

    Allows agents/orchestrators to lift the default ~4 MB message size
    ceiling when sending or receiving large payloads through the workflow
    runtime channel.
    """

    max_send_message_length: Optional[int] = None
    max_receive_message_length: Optional[int] = None
    keepalive_time_ms: Optional[int] = None
    keepalive_timeout_ms: Optional[int] = None

    def __post_init__(self) -> None:
        if (
            self.max_send_message_length is not None
            and self.max_send_message_length <= 0
        ):
            raise ValueError("max_send_message_length must be greater than 0")
        if (
            self.max_receive_message_length is not None
            and self.max_receive_message_length <= 0
        ):
            raise ValueError("max_receive_message_length must be greater than 0")
        if self.keepalive_time_ms is not None and self.keepalive_time_ms <= 0:
            raise ValueError("keepalive_time_ms must be greater than 0")
        if self.keepalive_timeout_ms is not None and self.keepalive_timeout_ms <= 0:
            raise ValueError("keepalive_timeout_ms must be greater than 0")


@dataclass
class AgentStateConfig:
    """
    State persistence configuration.

    Schema is auto-selected by agent/orchestrator type. Supply storage details
    and optional hooks; the framework injects the appropriate schema bundle.

    Examples:
        # Schema auto-selected by agent type
        config = AgentStateConfig(store=StateStoreService(...))
        agent = DurableAgent(state=config, ...)  # → AgentWorkflowState
        orch = DurableAgent(state=config, orchestration_mode='agent', ...)  # → LLMWorkflowState

        # With custom hooks
        config = AgentStateConfig(
            store=StateStoreService(...),
            entry_factory=custom_factory,
        )
    """

    store: "StateStoreService"
    default_state: Optional[Dict[str, Any] | BaseModel] = None
    state_key_prefix: Optional[str] = None

    # Hook overrides (optional - bundle provides defaults)
    entry_factory: Optional[EntryFactory] = None
    message_coercer: Optional[MessageCoercer] = None

    # Internal: schema bundle (injected by agent/orchestrator class)
    _state_model_bundle: Optional[StateModelBundle] = field(
        default=None, init=False, repr=False
    )

    def ensure_bundle(self, bundle: StateModelBundle) -> None:
        """
        Inject schema bundle (called by agent/orchestrator).

        Args:
            bundle: Schema bundle to use.

        Raises:
            RuntimeError: If different bundle already injected.
        """
        if self._state_model_bundle is not None:
            # Already set - verify it matches
            if (
                self._state_model_bundle.entry_model_cls != bundle.entry_model_cls
                or self._state_model_bundle.message_model_cls
                != bundle.message_model_cls
            ):
                raise RuntimeError(
                    f"State config already wired with "
                    f"{self._state_model_bundle.entry_model_cls.__name__} schema. "
                    f"Cannot inject {bundle.entry_model_cls.__name__} schema."
                )
            return  # Same bundle, no-op

        # Merge user hooks with bundle defaults
        self._state_model_bundle = StateModelBundle(
            entry_model_cls=bundle.entry_model_cls,
            message_model_cls=bundle.message_model_cls,
            entry_factory=self.entry_factory or bundle.entry_factory,
            message_coercer=self.message_coercer or bundle.message_coercer,
        )

    def get_state_model_bundle(self) -> StateModelBundle:
        """
        Get injected schema bundle.

        Returns:
            StateModelBundle with schema classes and hooks.

        Raises:
            RuntimeError: If bundle not injected yet.
        """
        if self._state_model_bundle is None:
            raise RuntimeError(
                "State config bundle not initialized. "
                "This should be injected by the agent/orchestrator class."
            )
        return self._state_model_bundle


@dataclass(frozen=True)
class ConfigFieldDescriptor:
    """Describes how a configuration key maps to an agent attribute.

    Attributes:
        target_type: Expected Python type for the coerced value.
        setter: Callable ``(agent, value) -> None`` that applies the value.
        sensitive: If ``True``, the value is redacted in log output.
        validator: Optional callable to validate/transform the coerced value.
        rebuilds_prompt: If ``True``, the prompt template is rebuilt after update.
    """

    target_type: Type
    setter: Callable[..., None]
    sensitive: bool = False
    validator: Optional[Callable[..., Any]] = None
    rebuilds_prompt: bool = False
    triggers_otel_reload: bool = False


# ---------------------------------------------------------------------------
# Built-in validators for ConfigFieldDescriptor
# ---------------------------------------------------------------------------

_config_logger = logging.getLogger(__name__)


def validate_non_empty_string(v: str) -> str:
    """Reject empty or whitespace-only strings."""
    if not v or not v.strip():
        raise ValueError("Value must not be empty")
    return v.strip()


def validate_max_iterations(v: int) -> int:
    """Ensure max_iterations is at least 1."""
    if v < 1:
        raise ValueError(f"max_iterations must be >= 1, got {v}")
    return v


def validate_tool_choice(v: str) -> str:
    """Warn if tool_choice is non-standard, but allow it."""
    allowed = {"auto", "none", "required"}
    if v.lower() not in allowed:
        _config_logger.warning(
            "tool_choice '%s' not in standard set %s; allowing anyway.", v, allowed
        )
    return v


def validate_otel_exporter_tracing(v: str) -> str:
    """Validate that the tracing exporter is a known AgentTracingExporter value."""
    try:
        AgentTracingExporter(v)
    except (ValueError, KeyError):
        raise ValueError(
            f"Unknown tracing exporter '{v}'. "
            f"Valid options: {[e.value for e in AgentTracingExporter]}"
        )
    return v


def validate_otel_exporter_logging(v: str) -> str:
    """Validate that the logging exporter is a known AgentLoggingExporter value."""
    try:
        AgentLoggingExporter(v)
    except (ValueError, KeyError):
        raise ValueError(
            f"Unknown logging exporter '{v}'. "
            f"Valid options: {[e.value for e in AgentLoggingExporter]}"
        )
    return v


@dataclass
class RuntimeSubscriptionConfig:
    """Configuration for subscribing to a Dapr Configuration Store at runtime.

    Attributes:
        store_name: Name of the Dapr configuration store component.
        default_key: Fallback key used when ``keys`` is empty (defaults to agent name).
        keys: Optional list of keys to subscribe to.
        metadata: Optional metadata for the configuration subscription.
        on_config_change: Optional callback invoked after each successful config update.
            Receives the normalized key and coerced value.
    """

    store_name: str
    default_key: Optional[str] = None
    keys: List[str] = field(default_factory=list)
    metadata: Dict[str, str] = field(default_factory=dict)
    on_config_change: Optional[Callable[[str, Any], None]] = None


@dataclass
class AgentRegistryConfig:
    """Configuration for agent registry storage."""

    store: StateStoreService
    team_name: Optional[str] = None


@dataclass
class AgentMemoryConfig:
    """Configuration wrapper for agent memory selection."""

    store: MemoryBase = field(default_factory=ConversationListMemory)


@dataclass
class AgentPubSubConfig:
    """Declarative pub/sub configuration for durable agents.

    Attributes:
        pubsub_name: Name of the Dapr pub/sub component to use for all agent traffic.
        agent_topic: Primary topic for direct messages to the agent. Defaults to ``name``.
        broadcast_topic: Optional topic shared by a team for broadcast messages.
    """

    pubsub_name: str
    agent_topic: Optional[str] = None
    broadcast_topic: Optional[str] = None


@dataclass
class PromptSection:
    """Reusable block for composing a structured system prompt."""

    title: str
    lines: List[str] = field(default_factory=list)

    def render(self, template_format: str) -> str:
        if not self.lines:
            return ""
        header = self.title.strip()
        body = "\n".join(f"- {line.strip()}" for line in self.lines if line.strip())
        section = f"{header}:\n{body}".strip()
        return (
            _ensure_jinja_placeholders(section)
            if template_format == "jinja2"
            else section
        )


@dataclass
class AgentProfileConfig:
    """
    High-level persona description for an agent.

    Mirrors common fields in OpenAI Agents SDK while remaining lightweight.
    """

    name: Optional[str] = None
    role: Optional[str] = None
    goal: Optional[str] = None
    instructions: List[str] = field(default_factory=list)
    style_guidelines: List[str] = field(default_factory=list)
    system_prompt: Optional[str] = None
    template_format: str = "jinja2"
    modules: Sequence[str] = field(default_factory=tuple)
    module_overrides: Dict[str, PromptSection] = field(default_factory=dict)


class ToolExecutionMode(StrEnum):
    """
    Enumeration of supported tool execution modes for durable agents.

    PARALLEL: All tool calls returned by the LLM in a single turn are executed
        concurrently via ``wf.when_all``. This is the default behaviour and
        provides the best latency when tools are independent.
    SEQUENTIAL: Tool calls are executed one after another in the order they
        were returned by the LLM. Use this when tools have side-effects that
        depend on the results of earlier calls in the same turn.
    """

    PARALLEL = "parallel"
    SEQUENTIAL = "sequential"


class OrchestrationMode(StrEnum):
    """
    Enumeration of supported orchestration strategies for durable agents.

    AGENT: Orchestration is driven by an LLM-generated plan that determines the next steps and agent interactions.
    RANDOM: Orchestration randomly selects agents or actions at each decision point, without a predetermined plan.
    ROUNDROBIN: Orchestration cycles through available agents or actions in a fixed order, ensuring equal opportunity for each participant.
    """

    AGENT = "agent"
    RANDOM = "random"
    ROUNDROBIN = "roundrobin"


@dataclass
class AgentExecutionConfig:
    """
    Dials to configure the agent execution.
    """

    # TODO: add a forceFinalAnswer field in case max_iterations is near/reached. Or do we have a conclusion baked in by default? Do we want this to derive a conclusion by default?
    # TODO: add stop_at_tokens
    max_iterations: int = 10
    tool_choice: Optional[str] = "auto"
    tool_execution_mode: ToolExecutionMode = ToolExecutionMode.PARALLEL
    orchestration_mode: Optional[OrchestrationMode] = None


@dataclass
class WorkflowRetryPolicy:
    """
    Configuration for durable retry policies in workflows.

    Attributes:
        max_attempts: Maximum number of retry attempts.
        initial_backoff_seconds: Initial backoff interval in seconds.
        max_backoff_seconds: Maximum backoff interval in seconds.
        backoff_multiplier: Multiplier for exponential backoff.
        retry_timeout: Optional total timeout for all retries in seconds.
    """

    max_attempts: Optional[int] = 3
    initial_backoff_seconds: Optional[int] = 5
    max_backoff_seconds: Optional[int] = 30
    backoff_multiplier: Optional[float] = 1.5
    retry_timeout: Optional[Union[int, None]] = None


class RuntimeConfigKey(StrEnum):
    """Supported keys for runtime configuration hot-reload.

    All profile keys use the ``agent_`` prefix to avoid ambiguity.
    Execution, LLM, and component keys are unprefixed.
    """

    # Profile fields
    AGENT_ROLE = "agent_role"
    AGENT_GOAL = "agent_goal"
    AGENT_INSTRUCTIONS = "agent_instructions"
    AGENT_SYSTEM_PROMPT = "agent_system_prompt"
    AGENT_STYLE_GUIDELINES = "agent_style_guidelines"

    # Execution fields
    MAX_ITERATIONS = "max_iterations"
    TOOL_CHOICE = "tool_choice"

    # LLM fields
    LLM_API_KEY = "llm_api_key"
    LLM_PROVIDER = "llm_provider"
    LLM_MODEL = "llm_model"

    # Component references
    AGENT_WORKFLOW = "agent_workflow"
    AGENT_REGISTRY = "agent_registry"
    AGENT_MEMORY = "agent_memory"

    # OTel fields — match standard env var names used throughout
    OTEL_SDK_DISABLED = "otel_sdk_disabled"
    OTEL_EXPORTER_OTLP_ENDPOINT = "otel_exporter_otlp_endpoint"
    OTEL_EXPORTER_OTLP_HEADERS = "otel_exporter_otlp_headers"
    OTEL_SERVICE_NAME = "otel_service_name"
    OTEL_TRACING_ENABLED = "otel_tracing_enabled"
    OTEL_TRACES_EXPORTER = "otel_traces_exporter"
    OTEL_LOGGING_ENABLED = "otel_logging_enabled"
    OTEL_LOGS_EXPORTER = "otel_logs_exporter"


class AgentTracingExporter(StrEnum):
    """
    Supported tracing exporters for Dapr Agents observability.
    """

    OTLP_GRPC = "otlp_grpc"
    OTLP_HTTP = "otlp_http"
    ZIPKIN = "zipkin"
    CONSOLE = "console"


class AgentLoggingExporter(StrEnum):
    """
    Supported logging exporters for Dapr Agents observability.
    """

    CONSOLE = "console"
    OTLP_GRPC = "otlp_grpc"
    OTLP_HTTP = "otlp_http"


@dataclass
class AgentObservabilityConfig:
    """
    Configuration settings for Dapr Agents observability features.

    Attributes:
        enabled: Enable/Disable observability.
        headers: Optional headers for observability exporters.
        auth_token: Optional authentication token for exporters.
        endpoint: Optional endpoint URL for observability exporters.
        service_name: Optional service name for observability data.
        logging_enabled: Enable/disable logging observability.
        logging_exporter: Logging exporter type.
        tracing_enabled: Enable/disable tracing observability.
        tracing_exporter: Tracing exporter type.
    """

    enabled: Optional[bool] = None
    headers: Dict[str, str] = field(default_factory=_empty_headers)
    auth_token: Optional[str] = None
    endpoint: Optional[str] = None
    service_name: Optional[str] = None
    logging_enabled: Optional[bool] = None
    logging_exporter: Optional[AgentLoggingExporter] = None
    tracing_enabled: Optional[bool] = None
    tracing_exporter: Optional[AgentTracingExporter] = None

    @classmethod
    def from_env(cls) -> "AgentObservabilityConfig":
        """Create observability config from standard OTEL environment variables.

        Uses standard OpenTelemetry env var names where available:
        - OTEL_SDK_DISABLED (inverted: disabled != "true" means enabled)
        - OTEL_EXPORTER_OTLP_ENDPOINT
        - OTEL_EXPORTER_OTLP_HEADERS (parses "Authorization=<token>" format)
        - OTEL_SERVICE_NAME
        - OTEL_TRACES_EXPORTER
        - OTEL_LOGS_EXPORTER
        - OTEL_TRACING_ENABLED (custom, no standard equivalent)
        - OTEL_LOGGING_ENABLED (custom, no standard equivalent)
        """
        headers: Dict[str, str] = {}
        raw_headers = getenv("OTEL_EXPORTER_OTLP_HEADERS")
        if raw_headers:
            for pair in raw_headers.split(","):
                pair = pair.strip()
                if "=" in pair:
                    k, v = pair.split("=", 1)
                    headers[k.strip()] = v.strip()

        logging_exporter: Optional[AgentLoggingExporter] = None
        if logging_exporter_str := getenv("OTEL_LOGS_EXPORTER"):
            try:
                logging_exporter = AgentLoggingExporter(logging_exporter_str)
            except (ValueError, KeyError):
                logging_exporter = AgentLoggingExporter.CONSOLE

        tracing_exporter: Optional[AgentTracingExporter] = None
        if tracing_exporter_str := getenv("OTEL_TRACES_EXPORTER"):
            try:
                tracing_exporter = AgentTracingExporter(tracing_exporter_str)
            except (ValueError, KeyError):
                tracing_exporter = AgentTracingExporter.CONSOLE

        enabled: Optional[bool] = None
        if getenv("OTEL_SDK_DISABLED") is not None:
            enabled = getenv("OTEL_SDK_DISABLED", "false").lower() != "true"

        logging_enabled: Optional[bool] = None
        if getenv("OTEL_LOGGING_ENABLED") is not None:
            logging_enabled = getenv("OTEL_LOGGING_ENABLED", "false").lower() == "true"

        tracing_enabled: Optional[bool] = None
        if getenv("OTEL_TRACING_ENABLED") is not None:
            tracing_enabled = getenv("OTEL_TRACING_ENABLED", "false").lower() == "true"

        return cls(
            enabled=enabled,
            headers=headers,
            endpoint=getenv("OTEL_EXPORTER_OTLP_ENDPOINT"),
            service_name=getenv("OTEL_SERVICE_NAME"),
            logging_enabled=logging_enabled,
            logging_exporter=logging_exporter,
            tracing_enabled=tracing_enabled,
            tracing_exporter=tracing_exporter,
        )


class AgentMetadata(BaseModel):
    """Metadata about an agent's configuration and capabilities."""

    appid: str = Field(
        ...,
        description="Dapr application ID (APP_ID) of the sidecar; may differ from the agent name",
    )
    type: str = Field(..., description="Type of the agent (e.g., standalone, durable)")
    orchestrator: bool = Field(
        False, description="Indicates if the agent is an orchestrator"
    )
    role: Optional[str] = Field(default=None, description="Role of the agent")
    goal: Optional[str] = Field(
        default=None, description="High-level objective of the agent"
    )
    instructions: Optional[List[str]] = Field(
        default=None, description="Instructions for the agent"
    )
    system_prompt: Optional[str] = Field(
        default=None, description="System prompt guiding the agent's behavior"
    )
    framework: Optional[str] = Field(
        default=None, description="Framework or library the agent is built with"
    )
    max_iterations: Optional[int] = Field(
        default=None, description="Maximum iterations for agent execution"
    )
    tool_choice: Optional[str] = Field(default=None, description="Tool choice strategy")
    metadata: Optional[Dict[str, Any]] = Field(
        default=None, description="Additional user-supplied metadata about the agent"
    )


class PubSubMetadata(BaseModel):
    """Pub/Sub configuration information."""

    resource_name: str = Field(..., description="Pub/Sub component name")
    broadcast_topic: Optional[str] = Field(
        default=None, description="Pub/Sub topic for broadcasting messages"
    )
    agent_topic: Optional[str] = Field(
        default=None, description="Pub/Sub topic for direct agent messages"
    )


class MemoryStoreMetadata(BaseModel):
    """Metadata about a single memory backing store."""

    type: str = Field(..., description="Implementation class name")
    resource_name: Optional[str] = Field(
        default=None, description="Dapr resource name for this store"
    )


class MemoryMetadata(BaseModel):
    """Memory configuration information."""

    short_term: Optional[MemoryStoreMetadata] = Field(
        default=None, description="Short-term workflow state store"
    )
    long_term: Optional[MemoryStoreMetadata] = Field(
        default=None, description="Long-term conversation memory store"
    )


class LLMMetadata(BaseModel):
    """LLM configuration information."""

    client: str = Field(..., description="LLM client used by the agent")
    provider: str = Field(..., description="LLM provider used by the agent")
    api: str = Field(default="unknown", description="API type used by the LLM client")
    model: str = Field(default="unknown", description="Model name or identifier")
    resource_name: Optional[str] = Field(
        default=None, description="Dapr resource name for the LLM client"
    )
    base_url: Optional[str] = Field(
        default=None, description="Base URL for the LLM API if applicable"
    )
    azure_endpoint: Optional[str] = Field(
        default=None, description="Azure endpoint if using Azure OpenAI"
    )
    azure_deployment: Optional[str] = Field(
        default=None, description="Azure deployment name if using Azure OpenAI"
    )
    prompt_template: Optional[str] = Field(
        default=None, description="Prompt template used by the agent"
    )


class ToolMetadata(BaseModel):
    """Metadata about a tool available to the agent."""

    name: str = Field(..., description="Name of the tool")
    description: str = Field(..., description="Description of the tool's functionality")
    args: str = Field(..., description="Arguments for the tool")


class RegistryMetadata(BaseModel):
    """Registry configuration information."""

    resource_name: Optional[str] = Field(
        None,
        description="Dapr resource name backing the registry (e.g. state store component)",
    )
    name: Optional[str] = Field(
        default=None, description="Logical team name the agent is registered under"
    )


class AgentMetadataSchema(BaseModel):
    """Schema for agent metadata including schema version."""

    version: str = Field(
        ...,
        description="Version of the schema used for the agent metadata.",
    )
    agent: AgentMetadata = Field(
        ..., description="Agent configuration and capabilities"
    )
    name: str = Field(
        ...,
        description="Logical agent name used as the registry key; distinct from agent.appid",
    )
    registered_at: str = Field(..., description="ISO 8601 timestamp of registration")
    pubsub: Optional[PubSubMetadata] = Field(
        None, description="Pub/sub configuration if enabled"
    )
    memory: Optional[MemoryMetadata] = Field(
        None, description="Memory configuration if enabled"
    )
    llm: Optional[LLMMetadata] = Field(None, description="LLM configuration")
    registry: Optional[RegistryMetadata] = Field(
        None, description="Registry configuration"
    )
    tools: Optional[List[ToolMetadata]] = Field(None, description="Available tools")

    @classmethod
    def export_json_schema(cls, version: str) -> Dict[str, Any]:
        """
        Export the JSON schema with version information.

        Args:
            version: The dapr-agents version for this schema

        Returns:
            JSON schema dictionary with metadata
        """
        schema = cls.model_json_schema()
        schema[_JSON_SCHEMA_KEY] = _JSON_SCHEMA_DRAFT_URL
        schema[_JSON_SCHEMA_VERSION_KEY] = version
        return schema
