"""
Dynamic Agent Execution Workflow for Temporal.

This is a DEDICATED workflow for dynamically created agents/teams.
It does NOT modify or depend on the existing agent_execution.py workflow.
"""

from dataclasses import dataclass, field
from datetime import timedelta
from typing import Optional, List, Dict, Any
from temporalio import workflow
from temporalio.common import RetryPolicy
import asyncio

with workflow.unsafe.imports_passed_through():
    from control_plane_api.app.activities.agent_activities import (
        execute_agent_llm,
        update_execution_status,
        ActivityExecuteAgentInput,
        ActivityUpdateExecutionInput,
    )


@dataclass
class DynamicAgentExecutionInput:
    """Input for dynamic agent/team execution workflow"""
    execution_id: str
    entity_id: str  # Ephemeral UUID (not in agents/teams table)
    entity_type: str  # "agent" or "team"
    organization_id: str
    prompt: str
    system_prompt: Optional[str] = None
    model_id: Optional[str] = None
    model_config: dict = None
    agent_config: dict = None
    mcp_servers: dict = None
    skill_ids: List[str] = None  # Skill IDs from Control Plane
    user_metadata: dict = None
    runtime_type: str = "claude_code"  # Default to claude_code for dynamic executions
    control_plane_url: Optional[str] = None
    api_key: Optional[str] = None
    initial_message_timestamp: Optional[str] = None
    graph_api_url: Optional[str] = None
    dataset_name: Optional[str] = None

    def __post_init__(self):
        if self.model_config is None:
            self.model_config = {}
        if self.agent_config is None:
            self.agent_config = {}
        if self.mcp_servers is None:
            self.mcp_servers = {}
        if self.skill_ids is None:
            self.skill_ids = []
        if self.user_metadata is None:
            self.user_metadata = {}


@dataclass
class DynamicAgentExecutionResult:
    """Result from dynamic agent/team execution"""
    status: str
    response: Optional[str] = None
    error_message: Optional[str] = None
    usage: dict = field(default_factory=dict)


@workflow.defn
class DynamicAgentExecutionWorkflow:
    """
    Dedicated workflow for dynamic agent/team execution.

    This workflow:
    1. Does NOT query agents/teams table (entities don't exist in DB)
    2. Uses entity_definition from execution_metadata
    3. Executes agent/team with Claude Code runtime
    4. Updates execution status
    """

    @workflow.run
    async def run(self, input: DynamicAgentExecutionInput) -> DynamicAgentExecutionResult:
        """Execute dynamic agent/team"""

        workflow.logger.info(
            "Starting dynamic agent execution",
            execution_id=input.execution_id,
            entity_type=input.entity_type,
            runtime=input.runtime_type,
        )

        # STEP 1: Update execution status to RUNNING
        await workflow.execute_activity(
            update_execution_status,
            ActivityUpdateExecutionInput(
                execution_id=input.execution_id,
                status="RUNNING",
                started_at=workflow.now().isoformat(),
            ),
            start_to_close_timeout=timedelta(seconds=10),
            retry_policy=RetryPolicy(
                maximum_attempts=3,
                initial_interval=timedelta(seconds=1),
            ),
        )

        try:
            # STEP 2: Execute agent with LLM
            llm_result = await workflow.execute_activity(
                execute_agent_llm,
                ActivityExecuteAgentInput(
                    execution_id=input.execution_id,
                    agent_id=input.entity_id,  # Use ephemeral entity_id
                    organization_id=input.organization_id,
                    prompt=input.prompt,
                    system_prompt=input.system_prompt,
                    model_id=input.model_id,
                    model_config=input.model_config,
                    mcp_servers=input.mcp_servers,
                    session_id=input.execution_id,  # Use execution_id as session_id
                    user_id=input.user_metadata.get("user_id"),
                    control_plane_url=input.control_plane_url,
                    api_key=input.api_key,
                    graph_api_url=input.graph_api_url,
                    dataset_name=input.dataset_name,
                ),
                start_to_close_timeout=timedelta(minutes=30),  # Long-running LLM execution
                heartbeat_timeout=timedelta(minutes=5),
                retry_policy=RetryPolicy(
                    maximum_attempts=1,  # No retries for LLM failures
                ),
            )

            # STEP 3: Update execution status to COMPLETED
            await workflow.execute_activity(
                update_execution_status,
                ActivityUpdateExecutionInput(
                    execution_id=input.execution_id,
                    status="COMPLETED",
                    completed_at=workflow.now().isoformat(),
                    response=llm_result.get("response"),
                    usage=llm_result.get("usage", {}),
                ),
                start_to_close_timeout=timedelta(seconds=10),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=1),
                ),
            )

            workflow.logger.info(
                "Dynamic execution completed successfully",
                execution_id=input.execution_id,
            )

            return DynamicAgentExecutionResult(
                status="COMPLETED",
                response=llm_result.get("response"),
                usage=llm_result.get("usage", {}),
            )

        except Exception as e:
            # STEP 4: Handle failure
            error_message = str(e)
            workflow.logger.error(
                "Dynamic execution failed",
                execution_id=input.execution_id,
                error=error_message,
            )

            await workflow.execute_activity(
                update_execution_status,
                ActivityUpdateExecutionInput(
                    execution_id=input.execution_id,
                    status="FAILED",
                    completed_at=workflow.now().isoformat(),
                    error_message=error_message,
                ),
                start_to_close_timeout=timedelta(seconds=10),
                retry_policy=RetryPolicy(
                    maximum_attempts=3,
                    initial_interval=timedelta(seconds=1),
                ),
            )

            return DynamicAgentExecutionResult(
                status="FAILED",
                error_message=error_message,
            )
