"""
ResilientAgent - Production-Ready Multi-Model Fallback for High Availability

Zero-config usage:
    from kaizen.agents import ResilientAgent

    agent = ResilientAgent(models=["gpt-4", "gpt-3.5-turbo", "local-model"])
    result = await agent.query("What is AI?")
    print(f"Used: {result['_fallback_strategy_used']}")

Progressive configuration:
    agent = ResilientAgent(
        models=["gpt-4", "gpt-3.5-turbo"],
        llm_provider="openai",
        temperature=0.7,
        max_tokens=300
    )

Environment variable support:
    KAIZEN_LLM_PROVIDER=openai
    KAIZEN_TEMPERATURE=0.7
    KAIZEN_MAX_TOKENS=300
"""

import os
from dataclasses import dataclass, field, replace
from typing import Any, Dict, List, Optional

from kaizen.core.base_agent import BaseAgent
from kaizen.signatures import InputField, OutputField, Signature
from kaizen.strategies.async_single_shot import AsyncSingleShotStrategy
from kaizen.strategies.fallback import FallbackStrategy

from kailash.nodes.base import NodeMetadata, register_node


class QuerySignature(Signature):
    """Signature for resilient query processing."""

    query: str = InputField(desc="Query to process")
    response: str = OutputField(desc="Response to query")


@dataclass
class ResilientConfig:
    """
    Configuration for Resilient Agent.

    All parameters have sensible defaults and can be overridden via:
    1. Constructor arguments (highest priority)
    2. Environment variables (KAIZEN_*)
    3. Default values (lowest priority)
    """

    # Fallback chain configuration
    models: List[str] = field(default_factory=lambda: ["gpt-4", "gpt-3.5-turbo"])

    # LLM configuration
    llm_provider: str = field(
        default_factory=lambda: os.getenv("KAIZEN_LLM_PROVIDER", "openai")
    )
    temperature: float = field(
        default_factory=lambda: float(os.getenv("KAIZEN_TEMPERATURE", "0.7"))
    )
    max_tokens: int = field(
        default_factory=lambda: int(os.getenv("KAIZEN_MAX_TOKENS", "300"))
    )

    # Technical configuration
    timeout: int = 30
    retry_attempts: int = 3
    provider_config: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        """Validate configuration."""
        if not self.models:
            raise ValueError(
                "ResilientConfig requires at least one model in fallback chain"
            )


@register_node()
class ResilientAgent(BaseAgent):
    """
    Production-ready Resilient Fallback Agent using FallbackStrategy.

    Features:
    - Zero-config with sensible defaults (GPT-4 → GPT-3.5)
    - Sequential fallback through model chain
    - Immediate return on first success
    - Track which strategy succeeded
    - Error summary for all failures
    - Built-in logging and error handling via BaseAgent

    Inherits from BaseAgent:
    - Signature-based query processing pattern
    - Fallback execution via FallbackStrategy
    - Error handling (ErrorHandlingMixin)
    - Performance tracking (PerformanceMixin)
    - Structured logging (LoggingMixin)

    Use Cases:
    - Multi-model fallback (GPT-4 → GPT-3.5 → local model)
    - Provider redundancy for high availability
    - Progressive degradation (primary → backup → emergency)
    - Cost optimization with quality tiers
    - Disaster recovery

    Performance:
    - Immediate return on first success
    - Error tracking for all failed attempts
    - No unnecessary retries
    - Typical latency: Primary model latency (if successful)

    Usage:
        # Zero-config (GPT-4 → GPT-3.5)
        import asyncio
        agent = ResilientAgent()

        result = asyncio.run(agent.query("What is Python?"))
        print(f"Response: {result['response']}")
        print(f"Used strategy: #{result['_fallback_strategy_used']}")

        # Custom fallback chain
        agent = ResilientAgent(
            models=["gpt-4", "gpt-3.5-turbo", "local-llama"]
        )

        # Error handling
        try:
            result = asyncio.run(agent.query("Complex query"))
        except Exception as e:
            errors = agent.get_error_summary()
            for error in errors:
                print(f"{error['strategy']}: {error['error']}")
    """

    # Node metadata for Studio discovery
    metadata = NodeMetadata(
        name="ResilientAgent",
        description="Multi-model fallback for high availability and progressive degradation",
        version="1.0.0",
        tags={
            "ai",
            "kaizen",
            "fallback",
            "resilient",
            "high-availability",
            "redundancy",
        },
    )

    def __init__(
        self,
        models: Optional[List[str]] = None,
        llm_provider: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        timeout: Optional[int] = None,
        retry_attempts: Optional[int] = None,
        provider_config: Optional[Dict[str, Any]] = None,
        config: Optional[ResilientConfig] = None,
        tool_registry: Optional["ToolRegistry"] = None,
        mcp_servers: Optional[List[Dict]] = None,
        **kwargs,
    ):
        """
        Initialize Resilient Agent with zero-config defaults.

        Args:
            models: Fallback chain of models (e.g., ["gpt-4", "gpt-3.5-turbo"])
            llm_provider: Override default LLM provider
            temperature: Override default temperature
            max_tokens: Override default max tokens
            timeout: Override default timeout
            retry_attempts: Override default retry attempts
            provider_config: Additional provider-specific configuration
            config: Full config object (overrides individual params)
            tool_registry: Optional tool registry for autonomous tool calling
            mcp_servers: Optional MCP server configurations for tool discovery
        """
        # If config object provided, use it; otherwise build from parameters
        if config is None:
            config = ResilientConfig()

            # Override defaults with provided parameters
            if models is not None:
                config = replace(config, models=models)
            if llm_provider is not None:
                config = replace(config, llm_provider=llm_provider)
            if temperature is not None:
                config = replace(config, temperature=temperature)
            if max_tokens is not None:
                config = replace(config, max_tokens=max_tokens)
            if timeout is not None:
                config = replace(config, timeout=timeout)
            if retry_attempts is not None:
                config = replace(config, retry_attempts=retry_attempts)
            if provider_config is not None:
                config = replace(config, provider_config=provider_config)

        # Merge timeout into provider_config
        if config.timeout and (
            not config.provider_config or "timeout" not in config.provider_config
        ):
            provider_cfg = (
                config.provider_config.copy() if config.provider_config else {}
            )
            provider_cfg["timeout"] = config.timeout
            config = replace(config, provider_config=provider_cfg)

        # Create strategies for each model in fallback chain
        strategies = []
        for model in config.models:
            # Each model gets its own strategy
            # In production, configure each with different models
            strategy = AsyncSingleShotStrategy()
            strategies.append(strategy)

        # Use FallbackStrategy with model strategies
        fallback_strategy = FallbackStrategy(strategies=strategies)

        # Initialize BaseAgent
        super().__init__(
            config=config,  # Auto-converted to BaseAgentConfig
            signature=QuerySignature(),
            strategy=fallback_strategy,
            tool_registry=tool_registry,
            mcp_servers=mcp_servers,
            **kwargs,
        )

        self.resilient_config = config

    async def query(self, query: str) -> Dict[str, Any]:
        """
        Process query with automatic fallback.

        Args:
            query: Query to process

        Returns:
            Dict[str, Any]: Result with response and fallback metadata

        Example:
            >>> import asyncio
            >>> agent = ResilientAgent()
            >>> result = asyncio.run(agent.query("What is Python?"))
            >>> print(result['response'])
            >>> print(f"Used strategy #{result['_fallback_strategy_used']}")
        """
        # Execute via strategy (handles fallback automatically)
        result = await self.strategy.execute(self, {"query": query})
        return result

    def query_sync(self, query: str) -> str:
        """
        Synchronous query (for compatibility).

        Args:
            query: Query to process

        Returns:
            str: Response text

        Example:
            >>> agent = ResilientAgent()
            >>> response = agent.query_sync("What is machine learning?")
        """
        import asyncio

        result = asyncio.run(self.query(query))
        return result.get("response", "No response")

    def get_error_summary(self) -> List[Dict[str, Any]]:
        """
        Get error summary from failed strategies.

        Returns:
            List[Dict[str, Any]]: Error summaries with strategy name, error message, and type

        Example:
            >>> agent = ResilientAgent()
            >>> try:
            ...     result = await agent.query("test")
            ... except Exception:
            ...     errors = agent.get_error_summary()
            ...     for error in errors:
            ...         print(f"{error['strategy']}: {error['error']}")
        """
        return self.strategy.get_error_summary()


# Convenience function for quick resilient queries
async def query_with_fallback(
    query: str,
    models: List[str] = None,
    llm_provider: str = "openai",
    temperature: float = 0.7,
) -> Dict[str, Any]:
    """
    Quick resilient query with default configuration.

    Args:
        query: Query to process
        models: Fallback chain of models
        llm_provider: LLM provider to use
        temperature: Temperature for generation

    Returns:
        Dict with response and fallback metadata

    Example:
        >>> import asyncio
        >>> from kaizen.agents.specialized.resilient import query_with_fallback
        >>>
        >>> result = asyncio.run(query_with_fallback(
        ...     "What is AI?",
        ...     models=["gpt-4", "gpt-3.5-turbo"]
        ... ))
        >>> print(f"Response: {result['response']}")
        >>> print(f"Used model #{result['_fallback_strategy_used']}")
    """
    if models is None:
        models = ["gpt-4", "gpt-3.5-turbo"]

    agent = ResilientAgent(
        models=models, llm_provider=llm_provider, temperature=temperature
    )

    return await agent.query(query)
