# Copyright 2025 Vijaykumar Singh <singhvjd@gmail.com>
#
# 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.

"""VerticalBase - Abstract base class for domain-specific assistants.

Implements the Template Method pattern where:
- VerticalBase defines the skeleton of creating an agent
- Concrete verticals override specific configuration steps

Design Patterns:
    - Template Method: Define skeleton, override steps
    - Factory Method: create_agent() factory
    - Registry: VerticalRegistry for discovery

Example:
    class MyCustomVertical(VerticalBase):
        name = "my_vertical"
        description = "Custom assistant for X"

        @classmethod
        def get_tools(cls) -> List[str]:
            return ["read", "write", "custom_tool"]

        @classmethod
        def get_system_prompt(cls) -> str:
            return "You are an expert in X..."
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Type

from victor.framework.tools import ToolSet


@dataclass
class StageDefinition:
    """Definition of a conversation stage for a vertical.

    Attributes:
        name: Stage name (e.g., "PLANNING", "EXECUTION")
        description: Human-readable description
        tools: Tools relevant to this stage
        keywords: Keywords that suggest this stage
        next_stages: Valid stages to transition to
        min_confidence: Minimum confidence to enter this stage
    """

    name: str
    description: str
    tools: Set[str] = field(default_factory=set)
    keywords: List[str] = field(default_factory=list)
    next_stages: Set[str] = field(default_factory=set)
    min_confidence: float = 0.5

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "tools": list(self.tools),
            "keywords": self.keywords,
            "next_stages": list(self.next_stages),
            "min_confidence": self.min_confidence,
        }


@dataclass
class VerticalConfig:
    """Configuration generated by a vertical.

    This is passed to Agent.create() to configure the agent.

    Attributes:
        tools: ToolSet configuration
        system_prompt: System prompt text
        stages: Stage definitions
        provider_hints: Hints for provider selection
        evaluation_criteria: Criteria for evaluating agent performance
        metadata: Additional vertical-specific metadata
    """

    tools: ToolSet
    system_prompt: str
    stages: Dict[str, StageDefinition] = field(default_factory=dict)
    provider_hints: Dict[str, Any] = field(default_factory=dict)
    evaluation_criteria: List[str] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_agent_kwargs(self) -> Dict[str, Any]:
        """Convert to kwargs for Agent.create().

        Returns:
            Dictionary of keyword arguments.
        """
        return {
            "tools": self.tools,
            # System prompt is typically handled separately
        }

    def to_dict(self) -> Dict[str, Any]:
        """Convert config to dictionary for dict-like access.

        Returns:
            Dictionary representation of the configuration.
        """
        return {
            "tools": self.tools.tools if self.tools else [],
            "system_prompt": self.system_prompt,
            "stages": {
                k: {"name": v.name, "description": v.description} for k, v in self.stages.items()
            },
            "provider_hints": self.provider_hints,
            "evaluation_criteria": self.evaluation_criteria,
            "metadata": self.metadata,
        }

    def keys(self):
        """Return keys for dict-like access compatibility.

        Returns:
            Keys iterator for dictionary representation.
        """
        return self.to_dict().keys()

    def __getitem__(self, key: str) -> Any:
        """Enable dict-like access: config["key"].

        Args:
            key: Configuration key

        Returns:
            Configuration value

        Raises:
            KeyError: If key not found
        """
        d = self.to_dict()
        if key not in d:
            raise KeyError(key)
        return d[key]

    def get(self, key: str, default: Any = None) -> Any:
        """Get config value with default, like dict.get().

        Args:
            key: Configuration key
            default: Default value if key not found

        Returns:
            Configuration value or default
        """
        try:
            return self[key]
        except KeyError:
            return default


class VerticalBase(ABC):
    """Abstract base class for domain-specific assistants.

    Implements the Template Method pattern. Subclasses must override:
    - name: Vertical identifier
    - description: Human-readable description
    - get_tools(): List of tool names
    - get_system_prompt(): System prompt text

    Optional overrides:
    - get_stages(): Stage definitions
    - get_provider_hints(): Provider selection hints
    - get_evaluation_criteria(): Performance criteria
    - customize_config(): Hook for final config customization

    Example:
        class SecurityAuditor(VerticalBase):
            name = "security_auditor"
            description = "Security vulnerability analysis"

            @classmethod
            def get_tools(cls):
                return ["read", "search", "security_scan"]

            @classmethod
            def get_system_prompt(cls):
                return "You are a security expert..."
    """

    # Subclasses must define these
    name: str = ""
    description: str = ""
    version: str = "1.0.0"

    @classmethod
    @abstractmethod
    def get_tools(cls) -> List[str]:
        """Get the list of tool names for this vertical.

        Returns:
            List of tool names to enable.
        """
        pass

    @classmethod
    @abstractmethod
    def get_system_prompt(cls) -> str:
        """Get the system prompt for this vertical.

        Returns:
            System prompt text with domain expertise.
        """
        pass

    @classmethod
    def get_stages(cls) -> Dict[str, StageDefinition]:
        """Get stage definitions for this vertical.

        Default implementation returns standard coding stages.
        Override for domain-specific stages.

        Returns:
            Dictionary mapping stage names to definitions.
        """
        return {
            "INITIAL": StageDefinition(
                name="INITIAL",
                description="Understanding the request",
                keywords=["what", "how", "explain", "help"],
                next_stages={"PLANNING", "READING"},
            ),
            "PLANNING": StageDefinition(
                name="PLANNING",
                description="Planning the approach",
                keywords=["plan", "approach", "strategy"],
                next_stages={"READING", "EXECUTION"},
            ),
            "READING": StageDefinition(
                name="READING",
                description="Gathering information",
                keywords=["read", "show", "find", "search"],
                next_stages={"ANALYSIS", "EXECUTION"},
            ),
            "ANALYSIS": StageDefinition(
                name="ANALYSIS",
                description="Analyzing information",
                keywords=["analyze", "review", "understand"],
                next_stages={"EXECUTION", "PLANNING"},
            ),
            "EXECUTION": StageDefinition(
                name="EXECUTION",
                description="Taking action",
                keywords=["change", "modify", "create", "implement"],
                next_stages={"VERIFICATION", "COMPLETION"},
            ),
            "VERIFICATION": StageDefinition(
                name="VERIFICATION",
                description="Verifying results",
                keywords=["test", "verify", "check", "validate"],
                next_stages={"COMPLETION", "EXECUTION"},
            ),
            "COMPLETION": StageDefinition(
                name="COMPLETION",
                description="Wrapping up",
                keywords=["done", "finished", "complete"],
                next_stages=set(),
            ),
        }

    @classmethod
    def get_provider_hints(cls) -> Dict[str, Any]:
        """Get hints for provider selection.

        Returns:
            Dictionary with provider preferences.
        """
        return {
            "preferred_providers": ["anthropic", "openai"],
            "min_context_window": 100000,
            "requires_tool_calling": True,
        }

    @classmethod
    def get_evaluation_criteria(cls) -> List[str]:
        """Get criteria for evaluating agent performance.

        Returns:
            List of evaluation criteria descriptions.
        """
        return [
            "Task completion accuracy",
            "Tool usage efficiency",
            "Response relevance",
        ]

    @classmethod
    def customize_config(cls, config: VerticalConfig) -> VerticalConfig:
        """Hook for final configuration customization.

        Override to modify the config before it's returned.

        Args:
            config: Configuration to customize.

        Returns:
            Customized configuration.
        """
        return config

    # =========================================================================
    # Extension Protocol Methods (Optional)
    # =========================================================================
    # These methods enable verticals to provide framework extensions.
    # Override them to integrate with the framework's middleware, safety,
    # prompt, and configuration systems.

    @classmethod
    def get_middleware(cls) -> List[Any]:
        """Get middleware implementations for this vertical.

        Override to provide vertical-specific middleware for tool
        execution processing.

        Returns:
            List of middleware implementations (MiddlewareProtocol)
        """
        return []

    @classmethod
    def get_safety_extension(cls) -> Optional[Any]:
        """Get safety extension for this vertical.

        Override to provide vertical-specific dangerous operation patterns.

        Returns:
            Safety extension (SafetyExtensionProtocol) or None
        """
        return None

    @classmethod
    def get_prompt_contributor(cls) -> Optional[Any]:
        """Get prompt contributor for this vertical.

        Override to provide vertical-specific task hints and prompt sections.

        Returns:
            Prompt contributor (PromptContributorProtocol) or None
        """
        return None

    @classmethod
    def get_mode_config_provider(cls) -> Optional[Any]:
        """Get mode configuration provider for this vertical.

        Override to provide vertical-specific operational modes.

        Returns:
            Mode config provider (ModeConfigProviderProtocol) or None
        """
        return None

    @classmethod
    def get_mode_config(cls) -> Dict[str, Any]:
        """Get mode configurations for this vertical.

        Returns operational modes like 'fast', 'thorough', 'explore' with
        their configurations (tool_budget, max_iterations, temperature).

        Default implementation provides standard modes. Override in subclasses
        for vertical-specific mode configurations.

        Returns:
            Dictionary mapping mode names to ModeConfig-like dicts.
        """
        return {
            "fast": {
                "name": "fast",
                "tool_budget": 10,
                "max_iterations": 20,
                "temperature": 0.7,
                "description": "Quick responses with limited tool usage",
            },
            "thorough": {
                "name": "thorough",
                "tool_budget": 50,
                "max_iterations": 50,
                "temperature": 0.7,
                "description": "Comprehensive analysis with extensive tool usage",
            },
            "explore": {
                "name": "explore",
                "tool_budget": 30,
                "max_iterations": 30,
                "temperature": 0.9,
                "description": "Exploratory mode with higher creativity",
            },
        }

    @classmethod
    def get_task_type_hints(cls) -> Dict[str, Any]:
        """Get task-type-specific prompt hints.

        Returns hints for common task types (edit, search, explain, etc.)
        with tool priorities and budget recommendations.

        Default implementation provides standard hints. Override in subclasses
        for vertical-specific task type hints.

        Returns:
            Dictionary mapping task types to TaskTypeHint-like dicts.
        """
        return {
            "edit": {
                "task_type": "edit",
                "hint": "[EDIT MODE] Read target files first, then make focused modifications.",
                "tool_budget": 15,
                "priority_tools": ["read", "edit", "grep"],
            },
            "search": {
                "task_type": "search",
                "hint": "[SEARCH MODE] Use semantic search and grep for efficient discovery.",
                "tool_budget": 10,
                "priority_tools": ["grep", "code_search", "ls"],
            },
            "explain": {
                "task_type": "explain",
                "hint": "[EXPLAIN MODE] Read relevant code and provide clear explanations.",
                "tool_budget": 8,
                "priority_tools": ["read", "grep", "overview"],
            },
            "debug": {
                "task_type": "debug",
                "hint": "[DEBUG MODE] Investigate systematically, check logs and error messages.",
                "tool_budget": 20,
                "priority_tools": ["read", "grep", "shell", "run_tests"],
            },
            "implement": {
                "task_type": "implement",
                "hint": "[IMPLEMENT MODE] Plan first, implement incrementally, verify each step.",
                "tool_budget": 30,
                "priority_tools": ["read", "write", "edit", "shell"],
            },
        }

    @classmethod
    def get_tool_dependency_provider(cls) -> Optional[Any]:
        """Get tool dependency provider for this vertical.

        Override to provide vertical-specific tool execution patterns.

        Returns:
            Tool dependency provider (ToolDependencyProviderProtocol) or None
        """
        return None

    @classmethod
    def get_tiered_tools(cls) -> Optional[Any]:
        """Get tiered tool configuration for intelligent selection.

        Override to provide vertical-specific tiered tool configuration
        for context-efficient tool selection. When implemented, this enables:

        1. Mandatory tools: Always included (e.g., read, ls)
        2. Vertical core: Always included for this vertical (e.g., web, fetch for research)
        3. Semantic pool: Selected based on query similarity and stage

        Example for research vertical:
            return TieredToolConfig(
                mandatory={"read", "ls"},
                vertical_core={"web", "fetch"},
                semantic_pool={"write", "edit", "grep", "search"},
                stage_tools={
                    "WRITING": {"write", "edit"},
                    "SEARCHING": {"web", "fetch", "grep"},
                },
                readonly_only_for_analysis=True,
            )

        Returns:
            TieredToolConfig or None (falls back to get_tools())
        """
        return None

    @classmethod
    def get_workflow_provider(cls) -> Optional[Any]:
        """Get workflow provider for this vertical.

        Override to provide vertical-specific workflows.

        Returns:
            Workflow provider (WorkflowProviderProtocol) or None
        """
        return None

    @classmethod
    def get_service_provider(cls) -> Optional[Any]:
        """Get service provider for this vertical.

        By default, returns a BaseVerticalServiceProvider that registers
        the vertical's prompt contributor, safety extension, mode config,
        and tool dependency providers with the DI container.

        Override to provide custom service registration logic.

        Returns:
            Service provider (ServiceProviderProtocol) or factory-created provider
        """
        try:
            from victor.verticals.base_service_provider import VerticalServiceProviderFactory

            return VerticalServiceProviderFactory.create(cls)
        except ImportError:
            return None

    @classmethod
    def get_extensions(cls) -> Any:
        """Get all extensions for this vertical.

        Aggregates all extension implementations for framework integration.
        Override for custom extension aggregation.

        Returns:
            VerticalExtensions or None
        """
        # Import here to avoid circular dependency
        try:
            from victor.verticals.protocols import VerticalExtensions

            safety = cls.get_safety_extension()
            prompt = cls.get_prompt_contributor()

            return VerticalExtensions(
                middleware=cls.get_middleware(),
                safety_extensions=[safety] if safety else [],
                prompt_contributors=[prompt] if prompt else [],
                mode_config_provider=cls.get_mode_config_provider(),
                tool_dependency_provider=cls.get_tool_dependency_provider(),
                workflow_provider=cls.get_workflow_provider(),
                service_provider=cls.get_service_provider(),
            )
        except ImportError:
            return None

    # =========================================================================
    # Template Method Implementation
    # =========================================================================

    @classmethod
    def get_config(cls) -> VerticalConfig:
        """Get the complete configuration for this vertical.

        This is the main template method that assembles the configuration
        by calling the various override points.

        Returns:
            Complete VerticalConfig for agent creation.
        """
        # Build tool set
        tool_names = cls.get_tools()
        tools = ToolSet.from_tools(tool_names)

        # Build config
        config = VerticalConfig(
            tools=tools,
            system_prompt=cls.get_system_prompt(),
            stages=cls.get_stages(),
            provider_hints=cls.get_provider_hints(),
            evaluation_criteria=cls.get_evaluation_criteria(),
            metadata={
                "vertical_name": cls.name,
                "vertical_version": cls.version,
                "description": cls.description,
            },
        )

        # Allow final customization
        return cls.customize_config(config)

    @classmethod
    def get_tool_set(cls) -> ToolSet:
        """Get the ToolSet for this vertical.

        Convenience method that returns just the tool configuration.

        Returns:
            Configured ToolSet.
        """
        return ToolSet.from_tools(cls.get_tools())

    @classmethod
    async def create_agent(
        cls,
        provider: str = "anthropic",
        model: Optional[str] = None,
        **kwargs: Any,
    ) -> Any:
        """Factory method to create an agent with this vertical's config.

        Args:
            provider: LLM provider name.
            model: Optional model identifier.
            **kwargs: Additional arguments passed to Agent.create().

        Returns:
            Configured Agent instance.
        """
        from victor.framework import Agent

        config = cls.get_config()
        agent_kwargs = config.to_agent_kwargs()
        agent_kwargs.update(kwargs)

        return await Agent.create(
            provider=provider,
            model=model,
            **agent_kwargs,
        )


class VerticalRegistry:
    """Registry for discovering and accessing verticals.

    Implements the Registry pattern for vertical discovery.

    Example:
        # Register a vertical
        VerticalRegistry.register(MyVertical)

        # List available verticals
        for name, vertical in VerticalRegistry.list_all():
            print(f"{name}: {vertical.description}")

        # Get a specific vertical
        coding = VerticalRegistry.get("coding")
    """

    _registry: Dict[str, Type[VerticalBase]] = {}

    @classmethod
    def register(cls, vertical: Type[VerticalBase]) -> None:
        """Register a vertical.

        Args:
            vertical: Vertical class to register.

        Raises:
            ValueError: If vertical has no name.
        """
        if not vertical.name:
            raise ValueError(f"Vertical {vertical.__name__} has no name defined")
        cls._registry[vertical.name] = vertical

    @classmethod
    def unregister(cls, name: str) -> None:
        """Unregister a vertical.

        Args:
            name: Vertical name to unregister.
        """
        if name in cls._registry:
            del cls._registry[name]

    @classmethod
    def get(cls, name: str) -> Optional[Type[VerticalBase]]:
        """Get a vertical by name.

        Args:
            name: Vertical name.

        Returns:
            Vertical class or None if not found.
        """
        return cls._registry.get(name)

    @classmethod
    def list_all(cls) -> List[tuple[str, Type[VerticalBase]]]:
        """List all registered verticals.

        Returns:
            List of (name, vertical_class) tuples.
        """
        return list(cls._registry.items())

    @classmethod
    def list_names(cls) -> List[str]:
        """List all registered vertical names.

        Returns:
            List of vertical names.
        """
        return list(cls._registry.keys())

    @classmethod
    def clear(cls) -> None:
        """Clear all registered verticals (for testing)."""
        cls._registry.clear()
