"""Runtime-based execution activities for Temporal workflows.

This module provides activities that use the RuntimeFactory/RuntimeRegistry system
for agent execution, supporting multiple runtimes (Agno/Default, Claude Code, etc.)
"""

from dataclasses import dataclass
from typing import Optional, List, Dict, Any
from temporalio import activity
import structlog
import os
import asyncio
import time

from control_plane_api.worker.runtimes.base import (
    RuntimeType,
    RuntimeExecutionContext,
    RuntimeExecutionResult,
)
from control_plane_api.worker.runtimes.factory import RuntimeFactory
from control_plane_api.worker.control_plane_client import get_control_plane_client
from control_plane_api.worker.services.cancellation_manager import CancellationManager
from control_plane_api.worker.services.runtime_analytics import submit_runtime_analytics
from control_plane_api.worker.services.analytics_service import AnalyticsService

logger = structlog.get_logger(__name__)


def inject_env_vars_into_mcp_servers(
    mcp_servers: Dict[str, Any],
    agent_config: Optional[Dict[str, Any]] = None,
    runtime_config: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Inject environment variables into MCP server configurations (runtime-agnostic).

    This ensures MCP servers have access to critical environment variables like:
    - KUBIYA_API_KEY: For API authentication
    - KUBIYA_API_BASE: For API base URL
    - Agent/team-specific environment variables from agent_config

    This function is runtime-agnostic and can be used by any runtime (Default, Claude Code, etc.)

    Args:
        mcp_servers: Dictionary of MCP server configurations
        agent_config: Optional agent configuration with env_vars
        runtime_config: Optional runtime configuration with env vars

    Returns:
        Modified MCP server configurations with injected env vars
    """
    if not mcp_servers:
        return mcp_servers

    # Collect environment variables to inject
    env_vars_to_inject = {}

    # Add Kubiya API credentials from OS environment
    kubiya_api_key = os.environ.get("KUBIYA_API_KEY")
    if kubiya_api_key:
        env_vars_to_inject["KUBIYA_API_KEY"] = kubiya_api_key

    kubiya_api_base = os.environ.get("KUBIYA_API_BASE")
    if kubiya_api_base:
        env_vars_to_inject["KUBIYA_API_BASE"] = kubiya_api_base

    # Add any env vars from agent_config
    if agent_config:
        agent_env_vars = agent_config.get("env_vars", {})
        if agent_env_vars and isinstance(agent_env_vars, dict):
            env_vars_to_inject.update(agent_env_vars)
        elif agent_env_vars:
            logger.warning("agent_config.env_vars is not a dict, skipping", type=type(agent_env_vars).__name__)

    # Also check runtime_config for env vars
    if runtime_config:
        runtime_env_vars = runtime_config.get("env", {})
        if runtime_env_vars and isinstance(runtime_env_vars, dict):
            env_vars_to_inject.update(runtime_env_vars)
        elif runtime_env_vars:
            logger.warning("runtime_config.env is not a dict, skipping", type=type(runtime_env_vars).__name__)

    if not env_vars_to_inject:
        logger.debug("No environment variables to inject into MCP servers")
        return mcp_servers

    logger.info(
        "Injecting environment variables into MCP servers",
        server_count=len(mcp_servers),
        env_var_keys=list(env_vars_to_inject.keys()),
    )

    def resolve_template_string(value: str, env_vars: Dict[str, str]) -> str:
        """
        Resolve template variables in a string.
        Replaces {{VAR_NAME}} with the value from env_vars.
        """
        import re
        result = value
        for var_name, var_value in env_vars.items():
            # Match {{VAR_NAME}} patterns
            pattern = r'\{\{' + re.escape(var_name) + r'\}\}'
            result = re.sub(pattern, var_value, result)
        return result

    def resolve_templates_in_dict(data: Any, env_vars: Dict[str, str]) -> Any:
        """
        Recursively resolve template variables in dictionaries, lists, and strings.
        """
        if isinstance(data, dict):
            return {k: resolve_templates_in_dict(v, env_vars) for k, v in data.items()}
        elif isinstance(data, list):
            return [resolve_templates_in_dict(item, env_vars) for item in data]
        elif isinstance(data, str):
            return resolve_template_string(data, env_vars)
        else:
            return data

    # Inject env vars into each MCP server
    modified_servers = {}
    for server_name, server_config in mcp_servers.items():
        try:
            # Handle different MCP server configuration formats
            if hasattr(server_config, 'env'):
                # StdioServerParameters or similar object with env attribute
                if server_config.env is None:
                    server_config.env = {}
                # Merge env vars (don't override existing ones from server config)
                server_config.env = {**env_vars_to_inject, **server_config.env}
                logger.debug(
                    f"Injected env vars into MCP server '{server_name}' (object with env attribute)",
                    env_count=len(server_config.env),
                )
            elif isinstance(server_config, dict):
                # Dictionary-based configuration
                # First, resolve template variables in the entire config
                server_config = resolve_templates_in_dict(server_config, env_vars_to_inject)

                # Then add env vars to the env field
                if 'env' not in server_config:
                    server_config['env'] = {}
                # Merge env vars (don't override existing ones from server config)
                server_config['env'] = {**env_vars_to_inject, **server_config['env']}
                logger.debug(
                    f"Injected env vars and resolved templates in MCP server '{server_name}' (dict config)",
                    env_count=len(server_config['env']),
                )
            else:
                # Unknown format - try to set env attribute directly
                try:
                    if not hasattr(server_config, 'env'):
                        setattr(server_config, 'env', {})
                    server_config.env = {**env_vars_to_inject, **getattr(server_config, 'env', {})}
                    logger.debug(
                        f"Injected env vars into MCP server '{server_name}' (setattr)",
                        env_count=len(server_config.env),
                    )
                except Exception as attr_error:
                    logger.warning(
                        f"Could not inject env vars into MCP server '{server_name}' - unsupported format",
                        server_type=type(server_config).__name__,
                        error=str(attr_error),
                    )

            modified_servers[server_name] = server_config

        except Exception as e:
            logger.error(
                f"Error injecting env vars into MCP server '{server_name}'",
                error=str(e),
                exc_info=True,
            )
            # Keep original server config if injection fails
            modified_servers[server_name] = server_config

    logger.info(
        "✅ Environment variables injected into MCP servers",
        server_count=len(modified_servers),
        env_vars_injected=list(env_vars_to_inject.keys()),
    )

    return modified_servers


@dataclass
class ActivityRuntimeExecuteInput:
    """Input for runtime-based execution activity"""
    execution_id: str
    agent_id: str
    organization_id: str
    prompt: str
    runtime_type: str = "default"  # "default", "claude_code", etc.
    system_prompt: Optional[str] = None
    model_id: Optional[str] = None
    model_config: Optional[Dict[str, Any]] = None
    agent_config: Optional[Dict[str, Any]] = None
    skills: Optional[List[Dict[str, Any]]] = None
    mcp_servers: Optional[Dict[str, Any]] = None
    conversation_history: Optional[List[Dict[str, Any]]] = None
    user_metadata: Optional[Dict[str, Any]] = None
    runtime_config: Optional[Dict[str, Any]] = None
    stream: bool = False
    conversation_turn: int = 1  # Track turn number for analytics

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


@activity.defn
async def execute_with_runtime(input: ActivityRuntimeExecuteInput) -> Dict[str, Any]:
    """
    Execute agent using the RuntimeFactory/RuntimeRegistry system.

    This activity:
    1. Creates a runtime based on runtime_type (default, claude_code, etc.)
    2. Builds execution context
    3. Executes (streaming or non-streaming)
    4. Returns results

    Args:
        input: Activity input with execution details and runtime_type

    Returns:
        Dict with response, usage, success flag, etc.
    """
    print("\n" + "="*80)
    print("🚀 RUNTIME-BASED EXECUTION START")
    print("="*80)
    print(f"Execution ID: {input.execution_id}")
    print(f"Agent ID: {input.agent_id}")
    print(f"Organization: {input.organization_id}")
    print(f"Runtime Type: {input.runtime_type}")
    print(f"Model: {input.model_id or 'default'}")
    print(f"Stream: {input.stream}")
    print(f"Skills: {len(input.skills)}")
    print(f"MCP Servers: {len(input.mcp_servers)}")
    print(f"Prompt: {input.prompt[:100]}..." if len(input.prompt) > 100 else f"Prompt: {input.prompt}")
    print("="*80 + "\n")

    activity.logger.info(
        "Executing with Runtime system",
        extra={
            "execution_id": input.execution_id,
            "agent_id": input.agent_id,
            "organization_id": input.organization_id,
            "runtime_type": input.runtime_type,
            "model_id": input.model_id,
            "stream": input.stream,
        }
    )

    try:
        # Track execution start time for analytics
        turn_start_time = time.time()

        # Get Control Plane client and cancellation manager
        control_plane = get_control_plane_client()
        cancellation_manager = CancellationManager()

        # Initialize analytics service for submission
        analytics_service = AnalyticsService(
            control_plane_url=control_plane.base_url if hasattr(control_plane, 'base_url') else "http://localhost:8000",
            api_key=os.environ.get("KUBIYA_API_KEY", ""),
        )

        # Parse runtime type
        try:
            runtime_type_enum = RuntimeType(input.runtime_type)
        except ValueError:
            logger.error(f"Invalid runtime_type: {input.runtime_type}, falling back to DEFAULT")
            runtime_type_enum = RuntimeType.DEFAULT

        # Create runtime using factory
        factory = RuntimeFactory()
        runtime = factory.create_runtime(
            runtime_type=runtime_type_enum,
            control_plane_client=control_plane,
            cancellation_manager=cancellation_manager,
        )

        logger.info(
            f"Created runtime",
            extra={
                "runtime_type": runtime_type_enum,
                "runtime_class": runtime.__class__.__name__,
                "capabilities": runtime.get_capabilities(),
            }
        )

        # Fetch and instantiate skills if runtime supports tools
        skills = input.skills or []
        if runtime.supports_tools():
            print(f"🔧 Fetching skills from Control Plane...")
            try:
                skill_configs = control_plane.get_skills(input.agent_id)
                if skill_configs:
                    print(f"✅ Resolved {len(skill_configs)} skill configs")
                    print(f"   Types: {[t.get('type') for t in skill_configs]}")
                    print(f"   Names: {[t.get('name') for t in skill_configs]}")
                    print(f"   Enabled: {[t.get('enabled', True) for t in skill_configs]}\n")

                    # DEBUG: Show full config for workflow_executor skills
                    for cfg in skill_configs:
                        if cfg.get('type') in ['workflow_executor', 'workflow']:
                            print(f"🔍 Workflow Executor Skill Config:")
                            print(f"   Name: {cfg.get('name')}")
                            print(f"   Type: {cfg.get('type')}")
                            print(f"   Enabled: {cfg.get('enabled', True)}")
                            print(f"   Config Keys: {list(cfg.get('configuration', {}).keys())}\n")

                    # Import here to avoid circular dependency
                    from control_plane_api.worker.services.skill_factory import SkillFactory

                    print(f"\n🔍 RUNTIME ACTIVITIES - BEFORE SKILL FACTORY:")
                    print(f"   input.execution_id: {input.execution_id}")
                    print(f"   type: {type(input.execution_id)}")
                    print(f"   bool(input.execution_id): {bool(input.execution_id)}")
                    print(f"   skill_configs count: {len(skill_configs)}\n")

                    # Determine runtime type from agent config or input
                    runtime_type = (input.agent_config or {}).get("runtime", input.runtime_type or "agno")

                    # For Claude Code runtime, pass skill configs directly
                    # (map_skills_to_tools in config.py handles dict mapping)
                    if runtime_type == "claude_code":
                        # Pass skill_configs directly - they will be mapped to tool names
                        # by map_skills_to_tools() in build_claude_options()
                        skills = skill_configs
                        print(f"✅ Prepared {len(skill_configs)} skill configs for Claude Code runtime")
                        print(f"   Skill types: {[s.get('type') for s in skill_configs]}\n")
                    else:
                        # For other runtimes (like Agno), use SkillFactory to instantiate Python skills
                        skill_factory = SkillFactory(runtime_type=runtime_type)
                        skill_factory.initialize()

                        skills = skill_factory.create_skills_from_list(
                            skill_configs,
                            execution_id=input.execution_id  # Pass execution_id for control plane streaming
                        )

                    if skills:
                        print(f"✅ Instantiated {len(skills)} skill(s)")
                        # Show types of instantiated skills
                        skill_types = [type(s).__name__ for s in skills]
                        print(f"   Skill classes: {skill_types}\n")
                    else:
                        print(f"⚠️  No skills were instantiated (all disabled or failed)\n")
                else:
                    print(f"⚠️  No skills found for agent\n")
            except Exception as e:
                print(f"❌ Error fetching skills: {str(e)}\n")
                logger.error("skill_fetch_error", error=str(e), exc_info=True)

        # Inject environment variables into MCP servers (runtime-agnostic)
        # This ensures all MCP servers have access to KUBIYA_API_KEY, KUBIYA_API_BASE, etc.
        mcp_servers_with_env = inject_env_vars_into_mcp_servers(
            mcp_servers=input.mcp_servers,
            agent_config=input.agent_config,
            runtime_config=input.runtime_config,
        )

        # Build execution context
        context = RuntimeExecutionContext(
            execution_id=input.execution_id,
            agent_id=input.agent_id,
            organization_id=input.organization_id,
            prompt=input.prompt,
            system_prompt=input.system_prompt,
            conversation_history=input.conversation_history,
            model_id=input.model_id,
            model_config=input.model_config,
            agent_config=input.agent_config,
            skills=skills,  # Use fetched skills
            mcp_servers=mcp_servers_with_env,  # Use MCP servers with injected env vars
            user_metadata=input.user_metadata,
            runtime_config=input.runtime_config,
        )

        # Execute based on streaming preference
        if input.stream:
            # Streaming execution
            logger.info("Starting streaming execution")
            accumulated_response = ""
            final_result = None

            # Generate unique message ID for this turn (execution_id + timestamp)
            message_id = f"{input.execution_id}_{int(time.time() * 1000000)}"

            # Track tool events published
            tool_events_published = {"start": 0, "complete": 0}

            # Define event callback for publishing tool events to Control Plane
            def event_callback(event: Dict):
                """Callback to publish events (tool start/complete, content chunks) to Control Plane SSE"""
                event_type = event.get("type")

                if event_type == "content_chunk":
                    # Content chunks are already handled below via result.response
                    pass
                elif event_type == "tool_start":
                    # Publish tool start event (synchronous - this runs in async context via callback)
                    try:
                        print(f"\n🔧 TOOL START EVENT: {event.get('tool_name')} (ID: {event.get('tool_execution_id')})")
                        control_plane.publish_event(
                            execution_id=input.execution_id,
                            event_type="tool_started",  # Match default runtime event type
                            data={
                                "tool_name": event.get("tool_name"),
                                "tool_execution_id": event.get("tool_execution_id"),
                                "tool_arguments": event.get("tool_args", {}),
                                "message": f"🔧 Executing tool: {event.get('tool_name')}",
                                "source": "agent",
                            }
                        )
                        tool_events_published["start"] += 1
                        print(f"📡 Published tool_started event #{tool_events_published['start']}: {event.get('tool_name')}")
                    except Exception as e:
                        logger.error(f"❌ Failed to publish tool_start event: {e}", exc_info=True)
                        print(f"❌ Failed to publish tool_start event: {e}")
                elif event_type == "tool_complete":
                    # Publish tool complete event
                    try:
                        status = event.get("status", "success")
                        icon = "✅" if status == "success" else "❌"
                        print(f"\n{icon} TOOL COMPLETE EVENT: {event.get('tool_name')} ({status})")
                        control_plane.publish_event(
                            execution_id=input.execution_id,
                            event_type="tool_completed",  # Match default runtime event type
                            data={
                                "tool_name": event.get("tool_name"),
                                "tool_execution_id": event.get("tool_execution_id"),
                                "status": status,
                                "tool_output": event.get("output"),
                                "tool_error": event.get("error"),
                                "message": f"{icon} Tool {status}: {event.get('tool_name')}",
                                "source": "agent",
                            }
                        )
                        tool_events_published["complete"] += 1
                        print(f"📡 Published tool_completed event #{tool_events_published['complete']}: {event.get('tool_name')}\n")
                    except Exception as e:
                        logger.error(f"❌ Failed to publish tool_complete event: {e}", exc_info=True)
                        print(f"❌ Failed to publish tool_complete event: {e}")

            # Stream execution with event callback
            # Note: AgnoRuntime publishes chunks via EventPublisher internally
            # But ClaudeCodeRuntime needs us to publish chunks here
            is_agno_runtime = runtime_type_enum == RuntimeType.DEFAULT

            async for result in runtime.stream_execute(context, event_callback):
                # Only process non-empty content (filter out empty strings and whitespace)
                if result.response and result.response.strip():
                    accumulated_response += result.response

                    # Publish chunks for non-Agno runtimes (e.g., Claude Code)
                    # AgnoRuntime publishes internally via EventPublisher to avoid duplicates
                    if not is_agno_runtime:
                        try:
                            await control_plane.publish_event_async(
                                execution_id=input.execution_id,
                                event_type="message_chunk",
                                data={
                                    "role": "assistant",
                                    "content": result.response,
                                    "is_chunk": True,
                                    "message_id": message_id,
                                }
                            )
                        except Exception as e:
                            logger.warning(f"Failed to publish streaming chunk: {e}")

                if result.finish_reason:
                    final_result = result
                    break

            if not final_result:
                raise RuntimeError("Streaming execution did not provide final result")

            # Log tool event summary
            print(f"\n📊 Tool Events Summary:")
            print(f"   tool_started events published: {tool_events_published['start']}")
            print(f"   tool_completed events published: {tool_events_published['complete']}")
            print(f"   tool_messages in result: {len(final_result.tool_messages or [])}\n")

            # Analytics now handled by separate Temporal activity in workflow
            # See: workflow calls submit_runtime_analytics_activity after this returns

            # Log before return to verify we reach this point
            logger.info(
                "activity_about_to_return_streaming",
                execution_id=input.execution_id,
                turn_number=input.conversation_turn,
                note="About to return streaming activity result to Temporal"
            )

            return {
                "success": final_result.success,
                "response": accumulated_response,
                "usage": final_result.usage or {},
                "model": final_result.model,
                "finish_reason": final_result.finish_reason,
                "tool_messages": final_result.tool_messages or [],
                "metadata": final_result.metadata or {},
                "error": final_result.error,
            }

        else:
            # Non-streaming execution
            logger.info("Starting non-streaming execution")
            result = await runtime.execute(context)

            # Analytics now handled by separate Temporal activity in workflow
            # See: workflow calls submit_runtime_analytics_activity after this returns

            # Log before return to verify we reach this point
            logger.info(
                "activity_about_to_return_non_streaming",
                execution_id=input.execution_id,
                turn_number=input.conversation_turn,
                note="About to return non-streaming activity result to Temporal"
            )

            return {
                "success": result.success,
                "response": result.response,
                "usage": result.usage or {},
                "model": result.model,
                "finish_reason": result.finish_reason,
                "tool_messages": result.tool_messages or [],
                "metadata": result.metadata or {},
                "error": result.error,
            }

    except Exception as e:
        logger.error(
            "Runtime execution failed",
            extra={
                "execution_id": input.execution_id,
                "runtime_type": input.runtime_type,
                "error": str(e),
            },
            exc_info=True,
        )

        # Publish error event to Control Plane for real-time UI updates
        try:
            from control_plane_api.worker.utils.error_publisher import (
                ErrorEventPublisher, ErrorSeverity, ErrorCategory
            )

            error_publisher = ErrorEventPublisher(control_plane)

            # Determine error category based on error message
            error_str = str(e).lower()
            category = ErrorCategory.UNKNOWN
            if "timeout" in error_str:
                category = ErrorCategory.TIMEOUT
            elif "import" in error_str or "module" in error_str:
                category = ErrorCategory.RUNTIME_INIT
            elif "api" in error_str or "model" in error_str or "anthropic" in error_str:
                category = ErrorCategory.MODEL_ERROR
            elif "network" in error_str or "connection" in error_str:
                category = ErrorCategory.NETWORK
            elif "auth" in error_str or "credential" in error_str:
                category = ErrorCategory.AUTHENTICATION

            await error_publisher.publish_error(
                execution_id=input.execution_id,
                exception=e,
                severity=ErrorSeverity.CRITICAL,
                category=category,
                stage="execution",
                component=f"{input.runtime_type}_runtime",
                operation="agent_execution",
                metadata={
                    "agent_id": input.agent_id,
                    "model_id": input.model_id,
                    "conversation_turn": input.conversation_turn,
                }
            )
        except Exception as publish_error:
            # Never let error publishing break the main flow
            logger.warning(
                "failed_to_publish_error_event",
                error=str(publish_error)
            )

        return {
            "success": False,
            "response": "",
            "usage": {},
            "model": input.model_id,
            "finish_reason": "error",
            "tool_messages": [],
            "metadata": {},
            "error": f"{type(e).__name__}: {str(e)}",
        }
