"""
Agentic Workflows - Pydantic schemas for Celery Beat driven enrichment.

This module defines the schema layer for the agentic workflow system that
enables background task enrichment through specialized agents.

Key Concepts:
    - WorkflowDefinition: Reusable workflow templates
    - WorkflowStep: Individual agent invocation within a workflow
    - PermissionRule: Controls auto-apply vs draft-for-approval
    - QualityGate: Criteria for task readiness

Usage:
    from lightwave.schema.pydantic.sst import (
        WorkflowDefinition,
        WorkflowStep,
        PermissionLevel,
        TaskMaturityState,
    )

    # Load workflow definition from SST
    workflow = WorkflowDefinition.from_sst("task_enrichment")

    # Check if agent can auto-apply
    if workflow.steps[0].permission == PermissionLevel.AUTO_APPLY:
        apply_directly()
    else:
        create_outbox_item()
"""

from __future__ import annotations

from enum import Enum
from functools import lru_cache
from typing import Any

import yaml
from pydantic import BaseModel, ConfigDict, Field

from lightwave.schema.core.paths import get_sst_path

# =============================================================================
# ENUMS
# =============================================================================


class TaskMaturityState(str, Enum):
    """Task lifecycle states from creation to execution-ready."""

    RAW = "raw"
    ENRICHING = "enriching"
    BLOCKED = "blocked"
    READY = "ready"
    IN_PROGRESS = "in_progress"
    REVIEW = "review"
    DONE = "done"


class PermissionLevel(str, Enum):
    """Permission levels for agent operations."""

    AUTO_APPLY = "auto_apply"
    DRAFT_FOR_APPROVAL = "draft_for_approval"
    SUGGEST_ONLY = "suggest_only"
    PROHIBITED = "prohibited"


class WorkflowTrigger(str, Enum):
    """What triggers a workflow to run."""

    CELERY_BEAT = "celery_beat"
    ON_CREATE = "on_create"
    ON_UPDATE = "on_update"
    MANUAL = "manual"
    WEBHOOK = "webhook"


# =============================================================================
# WORKFLOW STEP SCHEMAS
# =============================================================================


class StepOutputUpdate(BaseModel):
    """Defines how a step updates a field."""

    model_config = ConfigDict(extra="forbid")

    field: str = Field(..., description="Field to update")
    value: Any | None = Field(None, description="Static value to set")
    value_if_pass: Any | None = Field(None, description="Value if validation passes")
    value_if_fail: Any | None = Field(None, description="Value if validation fails")
    from_: str | None = Field(None, alias="from", description="Dynamic value from agent output")


class StepOutput(BaseModel):
    """Defines what a workflow step produces."""

    model_config = ConfigDict(extra="forbid")

    creates: str | None = Field(None, description="Model to create (e.g., ImplementationPlan)")
    updates: list[StepOutputUpdate] = Field(default_factory=list, description="Fields to update")


class StepInputFilter(BaseModel):
    """Defines which records a step should process."""

    model_config = ConfigDict(extra="allow")  # Allow arbitrary filter fields

    maturity_state: TaskMaturityState | None = None


class WorkflowStep(BaseModel):
    """A single step in a workflow - one agent invocation."""

    model_config = ConfigDict(extra="forbid")

    id: str = Field(..., description="Unique step identifier")
    agent: str = Field(..., description="Agent to invoke")
    description: str = Field(..., description="What this step does")
    input_filter: StepInputFilter | None = Field(None, description="Filter for input records")
    output: StepOutput | None = Field(None, description="What the step produces")
    permission: PermissionLevel = Field(
        PermissionLevel.DRAFT_FOR_APPROVAL,
        description="Permission level for this step",
    )
    retry_on_failure: bool = Field(False, description="Whether to retry on failure")
    max_retries: int = Field(3, description="Maximum retry attempts")
    depends_on: list[str] = Field(default_factory=list, description="Step IDs this depends on")
    frequency_override: str | None = Field(None, description="Override cron schedule for this step")


# =============================================================================
# WORKFLOW DEFINITION SCHEMA
# =============================================================================


class WorkflowDefinition(BaseModel):
    """A reusable workflow template defining an enrichment pipeline."""

    model_config = ConfigDict(extra="forbid")

    name: str = Field(..., description="Human-readable workflow name")
    description: str = Field(..., description="What this workflow does")
    trigger: WorkflowTrigger = Field(..., description="What triggers this workflow")
    schedule: str | None = Field(None, description="Cron schedule if trigger is celery_beat")
    steps: list[WorkflowStep] = Field(..., description="Ordered list of steps")
    enabled: bool = Field(True, description="Whether workflow is active")

    @classmethod
    @lru_cache(maxsize=1)
    def _load_sst(cls) -> dict[str, Any]:
        """Load the agentic-workflows.yaml SST definition."""
        from lightwave.schema.core.paths import get_sst_path

        yaml_path = get_sst_path("agentic_workflows")
        with open(yaml_path) as f:
            return yaml.safe_load(f)

    @classmethod
    def from_sst(cls, workflow_name: str) -> "WorkflowDefinition":
        """Load a workflow definition from SST."""
        data = cls._load_sst()
        workflows = data.get("workflow_definitions", {})

        if workflow_name not in workflows:
            available = list(workflows.keys())
            raise KeyError(f"Workflow '{workflow_name}' not found. Available: {available}")

        return cls.model_validate(workflows[workflow_name])

    @classmethod
    def get_all_from_sst(cls) -> list["WorkflowDefinition"]:
        """Load all workflow definitions from SST."""
        data = cls._load_sst()
        workflows = data.get("workflow_definitions", {})
        return [cls.model_validate(w) for w in workflows.values()]


# =============================================================================
# PERMISSION RULES SCHEMA
# =============================================================================


class OperationPermission(BaseModel):
    """Permission configuration for an operation type."""

    model_config = ConfigDict(extra="forbid")

    operation: str = Field(..., description="Operation name")
    default_level: PermissionLevel = Field(..., description="Default permission level")
    agent_overrides: dict[str, PermissionLevel] = Field(
        default_factory=dict,
        description="Agent-specific overrides",
    )

    def get_level_for_agent(self, agent_name: str) -> PermissionLevel:
        """Get the permission level for a specific agent."""
        return self.agent_overrides.get(agent_name, self.default_level)


class PermissionRules(BaseModel):
    """Complete permission configuration loaded from SST."""

    model_config = ConfigDict(extra="forbid")

    operation_defaults: dict[str, PermissionLevel]
    agent_overrides: dict[str, dict[str, PermissionLevel]]

    @classmethod
    @lru_cache(maxsize=1)
    def from_sst(cls) -> "PermissionRules":
        """Load permission rules from SST."""
        yaml_path = get_sst_path("agentic_workflows")
        with open(yaml_path) as f:
            data = yaml.safe_load(f)

        rules = data.get("permission_rules", {})
        return cls(
            operation_defaults={k: PermissionLevel(v) for k, v in rules.get("operation_defaults", {}).items()},
            agent_overrides={
                agent: {op: PermissionLevel(level) for op, level in overrides.items()}
                for agent, overrides in rules.get("agent_overrides", {}).items()
            },
        )

    def can_auto_apply(self, agent: str, operation: str) -> bool:
        """Check if an agent can auto-apply for an operation."""
        # Check agent override first
        if agent in self.agent_overrides:
            if operation in self.agent_overrides[agent]:
                return self.agent_overrides[agent][operation] == PermissionLevel.AUTO_APPLY

        # Fall back to default
        default = self.operation_defaults.get(operation, PermissionLevel.DRAFT_FOR_APPROVAL)
        return default == PermissionLevel.AUTO_APPLY

    def requires_approval(self, agent: str, operation: str) -> bool:
        """Check if an operation requires human approval."""
        return not self.can_auto_apply(agent, operation)


# =============================================================================
# QUALITY GATES SCHEMA
# =============================================================================


class QualityGateCriterion(BaseModel):
    """A single criterion for quality gate evaluation."""

    model_config = ConfigDict(extra="forbid")

    name: str
    required: bool = False
    score_weight: int = 0


class QualityGate(BaseModel):
    """Quality gate configuration for task readiness."""

    model_config = ConfigDict(extra="forbid")

    criteria: list[QualityGateCriterion]
    ready_threshold: int = Field(85, description="Score threshold for 'ready' state")
    enriching_threshold: int = Field(40, description="Score threshold for 'enriching' state")

    @classmethod
    @lru_cache(maxsize=1)
    def from_sst(cls) -> "QualityGate":
        """Load quality gate configuration from SST."""
        yaml_path = get_sst_path("agentic_workflows")
        with open(yaml_path) as f:
            data = yaml.safe_load(f)

        gates = data.get("quality_gates", {})
        task_criteria = gates.get("task_ready_criteria", {})
        scoring = task_criteria.get("scoring", {})

        criteria = []

        # Required criteria - can be list of single-key dicts or dict
        required_raw = task_criteria.get("required", [])
        if isinstance(required_raw, list):
            # List of single-key dicts: [{name: value}, ...]
            for item in required_raw:
                if isinstance(item, dict):
                    for name in item.keys():
                        score = scoring.get(name, 0)
                        criteria.append(QualityGateCriterion(name=name, required=True, score_weight=score))
        elif isinstance(required_raw, dict):
            for name in required_raw.keys():
                score = scoring.get(name, 0)
                criteria.append(QualityGateCriterion(name=name, required=True, score_weight=score))

        # Optional criteria - can be list of single-key dicts or dict
        optional_raw = task_criteria.get("optional", [])
        if isinstance(optional_raw, list):
            for item in optional_raw:
                if isinstance(item, dict):
                    for name in item.keys():
                        score = scoring.get(name, 0)
                        criteria.append(QualityGateCriterion(name=name, required=False, score_weight=score))
        elif isinstance(optional_raw, dict):
            for name in optional_raw.keys():
                score = scoring.get(name, 0)
                criteria.append(QualityGateCriterion(name=name, required=False, score_weight=score))

        thresholds = task_criteria.get("thresholds", {})
        return cls(
            criteria=criteria,
            ready_threshold=thresholds.get("ready", 85),
            enriching_threshold=thresholds.get("enriching", 40),
        )

    def evaluate(self, task_data: dict[str, Any]) -> tuple[TaskMaturityState, int]:
        """
        Evaluate a task against quality gates.

        Returns:
            Tuple of (recommended_state, score)
        """
        score = 0
        missing_required = []

        for criterion in self.criteria:
            # Check if criterion is met
            is_met = task_data.get(criterion.name, False)

            if is_met:
                score += criterion.score_weight
            elif criterion.required:
                missing_required.append(criterion.name)

        # Determine state
        if missing_required:
            return TaskMaturityState.BLOCKED, score
        elif score >= self.ready_threshold:
            return TaskMaturityState.READY, score
        elif score >= self.enriching_threshold:
            return TaskMaturityState.ENRICHING, score
        else:
            return TaskMaturityState.BLOCKED, score


# =============================================================================
# RETRY POLICY SCHEMA
# =============================================================================


class RetryPolicy(BaseModel):
    """Configuration for workflow retry behavior."""

    model_config = ConfigDict(extra="forbid")

    max_retries: int = Field(3, description="Maximum retry attempts")
    retry_delay_seconds: int = Field(300, description="Initial delay between retries")
    exponential_backoff: bool = Field(True, description="Use exponential backoff")
    backoff_multiplier: float = Field(2.0, description="Backoff multiplier")
    max_delay_seconds: int = Field(3600, description="Maximum delay cap")

    def get_delay(self, attempt: int) -> int:
        """Calculate delay for a given retry attempt."""
        if not self.exponential_backoff:
            return self.retry_delay_seconds

        delay = self.retry_delay_seconds * (self.backoff_multiplier**attempt)
        return min(int(delay), self.max_delay_seconds)

    @classmethod
    @lru_cache(maxsize=1)
    def from_sst(cls) -> "RetryPolicy":
        """Load default retry policy from SST."""
        yaml_path = get_sst_path("agentic_workflows")
        with open(yaml_path) as f:
            data = yaml.safe_load(f)

        policy = data.get("retry_policies", {}).get("default", {})
        return cls.model_validate(policy)


# =============================================================================
# AGENT REGISTRY SCHEMA
# =============================================================================


class AgentDefinition(BaseModel):
    """Definition of an available agent for workflows."""

    model_config = ConfigDict(extra="forbid")

    name: str
    source: str = Field(..., description="Path to agent definition or builtin identifier")
    status: str = Field("active", description="active, planned, deprecated")
    capabilities: list[str] = Field(default_factory=list)
    default_permission: PermissionLevel = Field(PermissionLevel.DRAFT_FOR_APPROVAL)

    @property
    def is_builtin(self) -> bool:
        """Check if this is a builtin agent."""
        return self.source.startswith("builtin:")

    @property
    def is_planned(self) -> bool:
        """Check if this agent is planned but not yet implemented."""
        return self.status == "planned"


class AgentRegistry(BaseModel):
    """Registry of all available agents."""

    model_config = ConfigDict(extra="forbid")

    agents: dict[str, AgentDefinition]

    @classmethod
    @lru_cache(maxsize=1)
    def from_sst(cls) -> "AgentRegistry":
        """Load agent registry from SST."""
        yaml_path = get_sst_path("agentic_workflows")
        with open(yaml_path) as f:
            data = yaml.safe_load(f)

        agents_data = data.get("agents", {})
        agents = {}
        for name, agent_def in agents_data.items():
            if name == "description":
                continue
            agents[name] = AgentDefinition(name=name, **agent_def)

        return cls(agents=agents)

    def get_agent(self, name: str) -> AgentDefinition:
        """Get an agent by name."""
        if name not in self.agents:
            available = list(self.agents.keys())
            raise KeyError(f"Agent '{name}' not found. Available: {available}")
        return self.agents[name]

    def get_active_agents(self) -> list[AgentDefinition]:
        """Get all active (non-planned) agents."""
        return [a for a in self.agents.values() if a.status == "active"]


# =============================================================================
# MODULE EXPORTS
# =============================================================================

__all__ = [
    # Enums
    "TaskMaturityState",
    "PermissionLevel",
    "WorkflowTrigger",
    # Workflow schemas
    "WorkflowDefinition",
    "WorkflowStep",
    "StepOutput",
    "StepOutputUpdate",
    "StepInputFilter",
    # Permission schemas
    "PermissionRules",
    "OperationPermission",
    # Quality gates
    "QualityGate",
    "QualityGateCriterion",
    # Retry policy
    "RetryPolicy",
    # Agent registry
    "AgentDefinition",
    "AgentRegistry",
]
