"""
Tool Registry data models.

Defines ToolDefinition, ExecutionResult, Decision, and lifecycle enums.
"""
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Awaitable, Callable, Optional, Union


# =============================================================================
# Enums
# =============================================================================

class ShadowMode(Enum):
    """Shadow execution modes for EMUs in shadow state."""
    SKIP = "skip"  # Don't execute shadow actions
    DRY_RUN = "dry_run"  # Validate inputs only
    EXECUTE_AND_COMPARE = "execute"  # Execute and compare with active


class LifecycleState(Enum):
    """EMU lifecycle states."""
    DRAFT = "draft"
    SHADOW = "shadow"
    CANARY = "canary"
    ACTIVE = "active"
    ARCHIVED = "archived"


class MissingToolBehavior(Enum):
    """Behavior when a tool referenced in an action is not registered."""
    ERROR = "error"  # ExecutionResult.success=False with error message
    SKIP = "skip"  # ExecutionResult.was_executed=False
    WARN = "warn"  # Same as skip but logs a warning


class RetryPolicy(str, Enum):
    """Retry policy for tool execution, matching EMU action ToolSpec."""
    NONE = "none"
    FIXED = "fixed"
    EXPONENTIAL = "exponential"


# =============================================================================
# Tool Definition
# =============================================================================

@dataclass
class ToolDefinition:
    """
    A registered tool with schema and handler.

    Attributes:
        name: Tool identifier (matches tool_id in EMU actions)
        description: Human-readable description
        input_schema: JSON Schema or simplified Python type schema
        projects: List of project IDs this tool belongs to (required)
        handler: Async function that executes the tool
        capabilities: List of capability tags (e.g., ["communication", "email"])
        constraints: List of constraint descriptions
        examples: List of example inputs/outputs
        idempotency_key: Template for generating idempotency keys
        rate_limit: Rate limit declaration (e.g., "100/hour")
        timeout_ms: Execution timeout in milliseconds
        signature_warnings: Validation warnings from handler signature check
        handler_pattern: Detected handler pattern ("args_dict" or "typed_params")
    """
    name: str
    description: str
    input_schema: dict[str, Any]
    projects: list[str]  # Required - list of project IDs this tool belongs to
    handler: Optional[Callable[[dict], Awaitable[dict[str, Any]]]] = None

    # AMI metadata
    capabilities: list[str] = field(default_factory=list)
    constraints: list[str] = field(default_factory=list)
    examples: list[dict] = field(default_factory=list)
    idempotency_key: Optional[str] = None
    rate_limit: Optional[str] = None
    timeout_ms: Optional[int] = None

    # Signature validation results (populated at registration time)
    signature_warnings: list[str] = field(default_factory=list)
    handler_pattern: Optional[str] = None  # "args_dict" or "typed_params"

    @property
    def has_handler(self) -> bool:
        """Check if a handler is bound to this tool."""
        return self.handler is not None

    def to_json_schema(self) -> dict[str, Any]:
        """Convert simplified schema to JSON Schema format."""
        if "type" in self.input_schema and self.input_schema.get("type") == "object":
            # Already in JSON Schema format
            return self.input_schema

        # Convert Python type hints to JSON Schema
        properties = {}
        required = []

        for key, type_hint in self.input_schema.items():
            prop_schema = self._type_to_json_schema(type_hint)
            properties[key] = prop_schema

            # Check if required (not Optional)
            if not self._is_optional(type_hint):
                required.append(key)

        return {
            "type": "object",
            "properties": properties,
            "required": required,
        }

    @staticmethod
    def _type_to_json_schema(type_hint: Any) -> dict[str, Any]:
        """Convert a Python type hint to JSON Schema."""
        # Handle None/NoneType
        if type_hint is None or type_hint is type(None):
            return {"type": "null"}

        # Handle basic types
        type_mapping = {
            str: {"type": "string"},
            int: {"type": "integer"},
            float: {"type": "number"},
            bool: {"type": "boolean"},
            list: {"type": "array"},
            dict: {"type": "object"},
        }

        if type_hint in type_mapping:
            return type_mapping[type_hint]

        # Handle typing module types
        origin = getattr(type_hint, "__origin__", None)
        args = getattr(type_hint, "__args__", ())

        # Handle Union types (including Optional)
        if origin is Union:
            # Check if it's Optional (Union[X, None])
            non_none_args = [a for a in args if a is not type(None)]
            if len(non_none_args) == 1:
                # Optional[X] - return schema for X
                return ToolDefinition._type_to_json_schema(non_none_args[0])
            else:
                # Multiple types - use anyOf
                return {
                    "anyOf": [ToolDefinition._type_to_json_schema(a) for a in non_none_args]
                }

        # Handle List[X]
        if origin is list:
            if args:
                return {
                    "type": "array",
                    "items": ToolDefinition._type_to_json_schema(args[0])
                }
            return {"type": "array"}

        # Handle Dict[K, V]
        if origin is dict:
            return {"type": "object"}

        # Fallback
        return {"type": "string"}

    @staticmethod
    def _is_optional(type_hint: Any) -> bool:
        """Check if a type hint is Optional."""
        origin = getattr(type_hint, "__origin__", None)
        args = getattr(type_hint, "__args__", ())

        if origin is Union:
            return type(None) in args
        return False


# =============================================================================
# Action Handlers Registry
# =============================================================================

@dataclass
class ActionHandlers:
    """Registry of handlers for different action types."""
    tool_call: bool = False
    context_directive: bool = False
    route: bool = False
    decision_prompt: bool = False

    # Actual handlers
    _route_handler: Optional[Callable] = field(default=None, repr=False)
    _decision_prompt_handler: Optional[Callable] = field(default=None, repr=False)
    _context_directive_handlers: dict = field(default_factory=dict, repr=False)


# =============================================================================
# Decision Model
# =============================================================================

@dataclass
class Decision:
    """
    An EMU decision with lifecycle context.

    Represents a selected EMU from invoke response with its action
    and lifecycle state for execution routing.
    """
    emu_key: str
    emu_version: int
    lifecycle_state: LifecycleState
    action: dict[str, Any]  # The action payload
    activation_id: Optional[str] = None
    canary_percentage: Optional[float] = None
    emu_id: Optional[str] = None
    score: Optional[float] = None
    resolved_idempotency_key: Optional[str] = None

    @classmethod
    def from_selected_item(cls, item: Any, lifecycle_state: str = "active") -> "Decision":
        """Create Decision from InvokeResponse.selected item."""
        # Convert string lifecycle state to enum
        state = LifecycleState(lifecycle_state.lower())

        # Extract action dict
        action_dict = {}
        if hasattr(item, "action") and item.action:
            action = item.action
            if hasattr(action, "__dict__"):
                action_dict = {
                    "type": getattr(action, "type", "unknown"),
                }
                if hasattr(action, "tool"):
                    action_dict["tool"] = action.tool
                if hasattr(action, "name"):
                    action_dict["name"] = action.name
                if hasattr(action, "args"):
                    action_dict["args"] = action.args
                if hasattr(action, "prompt"):
                    action_dict["prompt"] = action.prompt
                if hasattr(action, "route"):
                    action_dict["route"] = action.route
                if hasattr(action, "directive"):
                    action_dict["directive"] = action.directive
                if hasattr(action, "metadata"):
                    action_dict["metadata"] = action.metadata
            elif isinstance(action, dict):
                action_dict = action

        return cls(
            emu_key=item.emu_key,
            emu_version=item.emu_version or 1,
            lifecycle_state=state,
            action=action_dict,
            activation_id=item.activation_id,
            emu_id=item.id if hasattr(item, "id") else None,
            score=item.score if hasattr(item, "score") else None,
            resolved_idempotency_key=item.resolved_idempotency_key if hasattr(item, "resolved_idempotency_key") else None,
        )


# =============================================================================
# Execution Result
# =============================================================================

@dataclass
class ExecutionResult:
    """
    Result of executing an action.

    Attributes:
        decision: The decision that was executed
        success: Whether execution succeeded
        output: Output from the handler (if successful)
        error: Error message (if failed)
        duration_ms: Execution duration in milliseconds
        was_executed: Whether the action was actually executed (False for skipped)
        execution_mode: How the action was executed (live, shadow, canary, dry_run, skipped)
        resolved_idempotency_key: The resolved idempotency key (if template was provided)
    """
    decision: Decision
    success: bool
    output: Any = None
    error: Optional[str] = None
    duration_ms: int = 0
    was_executed: bool = True
    execution_mode: str = "live"  # live | shadow | canary | dry_run | skipped
    resolved_idempotency_key: Optional[str] = None
    side_effect: Optional[bool] = None  # Passthrough from EMU action ToolSpec


# =============================================================================
# ACK Result
# =============================================================================

@dataclass
class AckResult:
    """Result of acknowledging a decision."""
    activation_id: str
    success: bool
    error: Optional[str] = None
    cooldown_created: bool = False


# =============================================================================
# Combined Invoke and Execute Result
# =============================================================================

@dataclass
class InvokeAndExecuteResult:
    """Combined result from invoke_and_execute."""
    invoke_result: Any  # InvokeResponse
    execution_results: list[ExecutionResult]
    ack_results: list[AckResult] = field(default_factory=list)


# =============================================================================
# Executor Configuration
# =============================================================================

@dataclass
class ExecutorConfig:
    """Configuration for the ActionExecutor."""
    missing_tool_behavior: MissingToolBehavior = MissingToolBehavior.ERROR
    execution_timeout_ms: int = 30000  # 30 seconds default
    max_retries: int = 0
    retry_backoff_ms: int = 100


# =============================================================================
# Shadow Configuration
# =============================================================================

@dataclass
class ShadowConfig:
    """Configuration for shadow execution."""
    mode: ShadowMode = ShadowMode.SKIP
    comparison_handler: Optional[Callable] = None


# =============================================================================
# Canary Configuration
# =============================================================================

def make_sticky_decider(
    field: str = "user_id",
    salt: str = "",
) -> Callable[["Decision", dict], bool]:
    """
    Create a sticky session canary decider.

    Ensures the same user consistently sees the same version (canary or active)
    based on a hash of their identifier.

    Args:
        field: Context field to use for sticky session (default: "user_id")
        salt: Optional salt for the hash (useful for A/B test isolation)

    Returns:
        Decider function for use with configure_canary()

    Example:
        ```python
        executor.configure_canary(
            decider=make_sticky_decider("user_id")
        )

        # Or with salt for isolated experiments:
        executor.configure_canary(
            decider=make_sticky_decider("user_id", salt="experiment_v2")
        )
        ```
    """
    def decider(decision: "Decision", context: dict) -> bool:
        key = context.get(field, "")
        if not key:
            # No identifier - fall back to random
            import random
            return random.random() * 100 < (decision.canary_percentage or 0)

        # Create deterministic hash for sticky session
        hash_input = f"{salt}:{key}:{decision.emu_key}"
        bucket = hash(hash_input) % 100

        return bucket < (decision.canary_percentage or 0)

    return decider


@dataclass
class CanaryConfig:
    """Configuration for canary execution."""
    decider: Optional[Callable[[Decision, dict], bool]] = None
    metrics_handler: Optional[Callable] = None
    use_sticky_sessions: bool = False
    sticky_field: str = "user_id"

    def __post_init__(self):
        """Set up sticky decider if requested."""
        if self.use_sticky_sessions and self.decider is None:
            self.decider = make_sticky_decider(self.sticky_field)

    def should_execute_canary(self, decision: Decision, context: dict) -> bool:
        """Determine if canary should be executed."""
        if self.decider:
            return self.decider(decision, context)

        # Default: random based on percentage
        if decision.canary_percentage is None:
            return False

        import random
        return random.random() * 100 < decision.canary_percentage


# =============================================================================
# Executor Schema (for static analysis)
# =============================================================================

@dataclass
class ExecutorSchema:
    """Schema export format for static analysis."""
    version: str = "1.0.0"
    server_name: Optional[str] = None
    generated_at: Optional[str] = None
    tools: dict[str, dict] = field(default_factory=dict)
    action_handlers: dict[str, bool] = field(default_factory=dict)
    lifecycle: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "version": self.version,
            "server_name": self.server_name,
            "generated_at": self.generated_at or datetime.utcnow().isoformat() + "Z",
            "tools": self.tools,
            "action_handlers": self.action_handlers,
            "lifecycle": self.lifecycle,
        }


# =============================================================================
# Lifecycle Metrics (FR3.9)
# =============================================================================

@dataclass
class ExecutionMetricsSample:
    """A single execution metrics sample."""
    emu_key: str
    execution_mode: str  # "live" | "shadow" | "canary"
    success: bool
    duration_ms: int
    timestamp: datetime = field(default_factory=datetime.utcnow)
    error: Optional[str] = None
    output_hash: Optional[str] = None  # Hash of output for comparison


@dataclass
class LifecycleMetrics:
    """
    Aggregated metrics for lifecycle comparison.

    Provides statistics for comparing active vs canary or active vs shadow
    execution behavior.
    """
    emu_key: str
    window_start: datetime
    window_end: datetime

    # Active (live) metrics
    active_count: int = 0
    active_success_count: int = 0
    active_error_count: int = 0
    active_latency_sum_ms: int = 0
    active_latency_samples: list[int] = field(default_factory=list)

    # Canary metrics
    canary_count: int = 0
    canary_success_count: int = 0
    canary_error_count: int = 0
    canary_latency_sum_ms: int = 0
    canary_latency_samples: list[int] = field(default_factory=list)

    # Shadow metrics
    shadow_count: int = 0
    shadow_success_count: int = 0
    shadow_error_count: int = 0
    shadow_latency_sum_ms: int = 0
    shadow_latency_samples: list[int] = field(default_factory=list)

    # Comparison metrics
    output_match_count: int = 0
    output_mismatch_count: int = 0

    @property
    def active_success_rate(self) -> float:
        """Success rate for active executions."""
        if self.active_count == 0:
            return 0.0
        return self.active_success_count / self.active_count

    @property
    def canary_success_rate(self) -> float:
        """Success rate for canary executions."""
        if self.canary_count == 0:
            return 0.0
        return self.canary_success_count / self.canary_count

    @property
    def shadow_success_rate(self) -> float:
        """Success rate for shadow executions."""
        if self.shadow_count == 0:
            return 0.0
        return self.shadow_success_count / self.shadow_count

    @property
    def active_avg_latency_ms(self) -> float:
        """Average latency for active executions."""
        if self.active_count == 0:
            return 0.0
        return self.active_latency_sum_ms / self.active_count

    @property
    def canary_avg_latency_ms(self) -> float:
        """Average latency for canary executions."""
        if self.canary_count == 0:
            return 0.0
        return self.canary_latency_sum_ms / self.canary_count

    @property
    def shadow_avg_latency_ms(self) -> float:
        """Average latency for shadow executions."""
        if self.shadow_count == 0:
            return 0.0
        return self.shadow_latency_sum_ms / self.shadow_count

    @property
    def output_match_rate(self) -> float:
        """Rate of output matches in comparisons."""
        total = self.output_match_count + self.output_mismatch_count
        if total == 0:
            return 0.0
        return self.output_match_count / total

    def _percentile(self, samples: list[int], p: float) -> float:
        """Calculate percentile from samples."""
        if not samples:
            return 0.0
        sorted_samples = sorted(samples)
        idx = int(len(sorted_samples) * p / 100)
        idx = min(idx, len(sorted_samples) - 1)
        return float(sorted_samples[idx])

    @property
    def active_p50_latency_ms(self) -> float:
        """P50 latency for active executions."""
        return self._percentile(self.active_latency_samples, 50)

    @property
    def active_p99_latency_ms(self) -> float:
        """P99 latency for active executions."""
        return self._percentile(self.active_latency_samples, 99)

    @property
    def canary_p50_latency_ms(self) -> float:
        """P50 latency for canary executions."""
        return self._percentile(self.canary_latency_samples, 50)

    @property
    def canary_p99_latency_ms(self) -> float:
        """P99 latency for canary executions."""
        return self._percentile(self.canary_latency_samples, 99)

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization/export."""
        return {
            "emu_key": self.emu_key,
            "window_start": self.window_start.isoformat(),
            "window_end": self.window_end.isoformat(),
            "active": {
                "count": self.active_count,
                "success_rate": self.active_success_rate,
                "error_count": self.active_error_count,
                "avg_latency_ms": self.active_avg_latency_ms,
                "p50_latency_ms": self.active_p50_latency_ms,
                "p99_latency_ms": self.active_p99_latency_ms,
            },
            "canary": {
                "count": self.canary_count,
                "success_rate": self.canary_success_rate,
                "error_count": self.canary_error_count,
                "avg_latency_ms": self.canary_avg_latency_ms,
                "p50_latency_ms": self.canary_p50_latency_ms,
                "p99_latency_ms": self.canary_p99_latency_ms,
            },
            "shadow": {
                "count": self.shadow_count,
                "success_rate": self.shadow_success_rate,
                "error_count": self.shadow_error_count,
                "avg_latency_ms": self.shadow_avg_latency_ms,
            },
            "comparison": {
                "output_match_rate": self.output_match_rate,
                "output_match_count": self.output_match_count,
                "output_mismatch_count": self.output_mismatch_count,
            },
        }


@dataclass
class ShadowComparisonResult:
    """
    Result of comparing active and shadow executions.

    Used when shadow mode is EXECUTE_AND_COMPARE and both active and shadow
    versions of an EMU are executed.
    """
    emu_key: str
    active_result: Optional["ExecutionResult"] = None
    shadow_result: Optional["ExecutionResult"] = None
    outputs_match: bool = False
    latency_diff_ms: int = 0
    both_succeeded: bool = False
    comparison_error: Optional[str] = None

    @classmethod
    def compare(
        cls,
        emu_key: str,
        active_result: Optional["ExecutionResult"],
        shadow_result: Optional["ExecutionResult"],
    ) -> "ShadowComparisonResult":
        """
        Create a comparison result from active and shadow execution results.

        Args:
            emu_key: The EMU key being compared
            active_result: Result from active execution (may be None)
            shadow_result: Result from shadow execution (may be None)

        Returns:
            ShadowComparisonResult with comparison details
        """
        if active_result is None or shadow_result is None:
            return cls(
                emu_key=emu_key,
                active_result=active_result,
                shadow_result=shadow_result,
                comparison_error="Missing result for comparison",
            )

        # Check if both succeeded
        both_succeeded = active_result.success and shadow_result.success

        # Compare outputs
        outputs_match = False
        if both_succeeded:
            # Simple equality check - could be enhanced with deep comparison
            outputs_match = active_result.output == shadow_result.output

        # Calculate latency difference
        latency_diff_ms = shadow_result.duration_ms - active_result.duration_ms

        return cls(
            emu_key=emu_key,
            active_result=active_result,
            shadow_result=shadow_result,
            outputs_match=outputs_match,
            latency_diff_ms=latency_diff_ms,
            both_succeeded=both_succeeded,
        )
