"""
Refactored Agent executor service using runtime abstraction.

This version delegates execution to pluggable runtime implementations,
making the code more maintainable and extensible.
"""

from typing import Dict, Any, Optional
import structlog
import time
import os
from datetime import datetime, timezone

from control_plane_api.worker.control_plane_client import ControlPlaneClient
from control_plane_api.worker.services.session_service import SessionService
from control_plane_api.worker.services.cancellation_manager import CancellationManager
from control_plane_api.worker.services.analytics_service import AnalyticsService
from control_plane_api.worker.services.runtime_analytics import submit_runtime_analytics
from control_plane_api.worker.runtimes import (
    RuntimeFactory,
    RuntimeType,
    RuntimeExecutionContext,
)
from control_plane_api.worker.utils.streaming_utils import StreamingHelper

logger = structlog.get_logger()


class AgentExecutorServiceV2:
    """
    Service for executing agents using runtime abstraction.

    This service orchestrates agent execution by:
    1. Loading session history
    2. Selecting appropriate runtime based on agent config
    3. Delegating execution to the runtime
    4. Persisting session after execution
    """

    def __init__(
        self,
        control_plane: ControlPlaneClient,
        session_service: SessionService,
        cancellation_manager: CancellationManager,
    ):
        """
        Initialize the agent executor service.

        Args:
            control_plane: Control Plane API client
            session_service: Session management service
            cancellation_manager: Execution cancellation manager
        """
        self.control_plane = control_plane
        self.session_service = session_service
        self.cancellation_manager = cancellation_manager
        self.runtime_factory = RuntimeFactory()

        # Initialize analytics service for tracking LLM usage, tool calls, etc.
        control_plane_url = os.getenv("CONTROL_PLANE_URL", "http://localhost:8000")
        api_key = os.getenv("KUBIYA_API_KEY", "")
        self.analytics_service = AnalyticsService(control_plane_url, api_key)

    async def execute(self, input: Any) -> Dict[str, Any]:
        """
        Execute an agent using the configured runtime.

        This method:
        1. Loads session history
        2. Determines runtime type from agent config
        3. Creates runtime instance
        4. Executes agent via runtime
        5. Persists session
        6. Returns standardized result

        Args:
            input: AgentExecutionInput with execution details

        Returns:
            Dict with response, usage, success flag, runtime_type, etc.
        """
        execution_id = input.execution_id

        print("\n" + "=" * 80)
        print("🤖 AGENT EXECUTION START (Runtime-Abstracted)")
        print("=" * 80)
        print(f"Execution ID: {execution_id}")
        print(f"Agent ID: {input.agent_id}")
        print(f"Organization: {input.organization_id}")
        print(f"Model: {input.model_id or 'default'}")
        print(f"Session ID: {input.session_id}")
        print(
            f"Prompt: {input.prompt[:100]}..."
            if len(input.prompt) > 100
            else f"Prompt: {input.prompt}"
        )
        print("=" * 80 + "\n")

        logger.info(
            "agent_execution_start",
            execution_id=execution_id[:8],
            agent_id=input.agent_id,
            session_id=input.session_id,
        )

        try:
            # Capture timestamp at start of execution for accurate user message timestamp
            from datetime import datetime, timezone
            user_message_timestamp = datetime.now(timezone.utc).isoformat()

            # STEP 1: Load session history
            print("📚 Loading session history...")
            session_history = self.session_service.load_session(
                execution_id=execution_id, session_id=input.session_id
            )

            if session_history:
                print(
                    f"✅ Loaded {len(session_history)} messages from previous session\n"
                )
            else:
                print("ℹ️  Starting new conversation\n")

            # STEP 2: Determine runtime type
            agent_config = input.agent_config or {}
            runtime_type_str = agent_config.get("runtime", "default")
            runtime_type = self.runtime_factory.parse_runtime_type(runtime_type_str)

            print(f"🔌 Runtime Type: {runtime_type.value}")
            print(f"   Framework: {self._get_framework_name(runtime_type)}\n")

            logger.info(
                "runtime_selected",
                execution_id=execution_id[:8],
                runtime=runtime_type.value,
            )

            # STEP 3: Create runtime instance
            print(f"⚙️  Creating runtime instance...")
            runtime = self.runtime_factory.create_runtime(
                runtime_type=runtime_type,
                control_plane_client=self.control_plane,
                cancellation_manager=self.cancellation_manager,
            )
            print(f"✅ Runtime created: {runtime.get_runtime_info()}\n")

            # STEP 4: Get skills (if runtime supports tools)
            skills = []
            if runtime.supports_tools():
                print(f"🔧 Fetching skills from Control Plane...")
                try:
                    skill_configs = self.control_plane.get_skills(input.agent_id)
                    if skill_configs:
                        print(f"✅ Resolved {len(skill_configs)} skills")
                        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 services.skill_factory import SkillFactory

                        # Create factory instance for the current runtime
                        skill_factory = SkillFactory(runtime_type=runtime_type.value)
                        skill_factory.initialize()

                        skills = skill_factory.create_skills_from_list(
                            skill_configs, execution_id=execution_id
                        )

                        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)

            # STEP 5: Inject environment variables into MCP servers (runtime-agnostic)
            print("🔐 Injecting environment variables into MCP servers...")
            from control_plane_api.worker.activities.runtime_activities import inject_env_vars_into_mcp_servers
            mcp_servers_with_env = inject_env_vars_into_mcp_servers(
                mcp_servers=input.mcp_servers,
                agent_config=agent_config,
                runtime_config=agent_config.get("runtime_config"),
            )

            # STEP 6: Build execution context
            print("📦 Building execution context...")
            context = RuntimeExecutionContext(
                execution_id=execution_id,
                agent_id=input.agent_id,
                organization_id=input.organization_id,
                prompt=input.prompt,
                system_prompt=input.system_prompt,
                conversation_history=session_history,
                model_id=input.model_id,
                model_config=input.model_config,
                agent_config=agent_config,
                skills=skills,
                mcp_servers=mcp_servers_with_env,  # Use MCP servers with injected env vars
                user_metadata=input.user_metadata,
                runtime_config=agent_config.get("runtime_config"),
            )
            print("✅ Context ready\n")

            # STEP 7: Publish user message to stream before execution
            # This ensures chronological ordering in UI
            turn_number = len(session_history) // 2 + 1
            user_message_id = f"{execution_id}_user_{turn_number}"
            self.control_plane.publish_event(
                execution_id=execution_id,
                event_type="message",
                data={
                    "role": "user",
                    "content": input.prompt,
                    "timestamp": user_message_timestamp,
                    "message_id": user_message_id,
                    "user_id": input.user_id,
                    "user_name": getattr(input, "user_name", None),
                    "user_email": getattr(input, "user_email", None),
                    "user_avatar": getattr(input, "user_avatar", None),
                }
            )
            print(f"   📤 Published user message to stream (ID: {user_message_id})\n")

            # Publish status message chunk: Connecting to worker
            assistant_message_id = f"{execution_id}_assistant_{turn_number}"
            self.control_plane.publish_event_async(
                execution_id=execution_id,
                event_type="message_chunk",
                data={
                    "role": "assistant",
                    "content": "Connecting to worker...\n",
                    "is_chunk": True,
                    "message_id": assistant_message_id,
                }
            )
            print(f"   📤 Published status: 'Connecting to worker...'\n")

            # STEP 8: Execute via runtime (with streaming if supported)
            print("⚡ Executing via runtime...\n")

            # Track turn start time for analytics
            turn_start_time = time.time()

            # Create streaming helper for tracking tool messages (used in both streaming and non-streaming)
            streaming_helper = StreamingHelper(
                control_plane_client=self.control_plane, execution_id=execution_id
            )

            if runtime.supports_streaming():
                result = await self._execute_streaming(runtime, context, input, streaming_helper)
            else:
                result = await runtime.execute(context)

            # Track turn end time
            turn_end_time = time.time()

            print("\n✅ Runtime execution completed!")
            print(f"   Response Length: {len(result.response)} chars")
            print(f"   Success: {result.success}\n")

            logger.info(
                "agent_execution_completed",
                execution_id=execution_id[:8],
                success=result.success,
                response_length=len(result.response),
            )

            # STEP 7.5: Submit analytics (non-blocking, fire-and-forget)
            if result.success and result.usage:
                try:
                    # Submit analytics in the background (doesn't block execution)
                    await submit_runtime_analytics(
                        result=result,
                        execution_id=execution_id,
                        turn_number=turn_number,
                        turn_start_time=turn_start_time,
                        turn_end_time=turn_end_time,
                        analytics_service=self.analytics_service,
                    )
                    logger.info(
                        "analytics_submitted",
                        execution_id=execution_id[:8],
                        tokens=result.usage.get("total_tokens", 0),
                    )
                except Exception as analytics_error:
                    # Analytics failures should not break execution
                    logger.warning(
                        "analytics_submission_failed",
                        execution_id=execution_id[:8],
                        error=str(analytics_error),
                    )

            # STEP 7: Persist session
            if result.success and result.response:
                print("💾 Persisting session history...")

                # Finalize streaming to transition to post-tool phase
                streaming_helper.finalize_streaming()

                # Build user message
                user_message = {
                    "role": "user",
                    "content": input.prompt,
                    "timestamp": user_message_timestamp,  # Use timestamp from start
                    "message_id": f"{execution_id}_user_{turn_number}",
                    "user_id": input.user_id,
                    "user_name": getattr(input, "user_name", None),
                    "user_email": getattr(input, "user_email", None),
                }

                # Build structured messages using StreamingHelper
                # This properly splits assistant messages around tool usage
                new_messages = streaming_helper.build_structured_messages(
                    execution_id=execution_id,
                    turn_number=turn_number,
                    user_message=user_message,
                )

                print(f"   📊 Built {len(new_messages)} structured messages for this turn")
                if streaming_helper.has_any_tools:
                    print(f"   🔧 Tools were used - assistant message split into phases")
                    assistant_parts = streaming_helper.get_assistant_message_parts()
                    for part in assistant_parts:
                        print(f"      - {part['phase']}: {len(part['content'])} chars")

                # Combine with previous history
                complete_session = session_history + new_messages

                # CRITICAL: Deduplicate messages by message_id to prevent duplicates
                # Use a dict to track seen message_ids, keeping the first occurrence
                seen_message_ids = {}
                deduplicated_session = []
                for msg in complete_session:
                    msg_id = msg.get("message_id")
                    if msg_id and msg_id in seen_message_ids:
                        # Skip duplicate message
                        continue
                    if msg_id:
                        seen_message_ids[msg_id] = True
                    deduplicated_session.append(msg)

                # CRITICAL: Sort by timestamp to ensure chronological order
                # Tool messages happen DURING streaming, so they need to be interleaved with user/assistant messages
                deduplicated_session.sort(key=lambda msg: msg.get("timestamp", ""))
                print(f"   ✅ Messages deduplicated ({len(complete_session)} -> {len(deduplicated_session)}) and sorted by timestamp")

                complete_session = deduplicated_session

                success = self.session_service.persist_session(
                    execution_id=execution_id,
                    session_id=input.session_id or execution_id,
                    user_id=input.user_id,
                    messages=complete_session,
                    metadata={
                        "agent_id": input.agent_id,
                        "organization_id": input.organization_id,
                        "runtime_type": runtime_type.value,
                        "turn_count": len(complete_session),
                    },
                )

                if success:
                    print(
                        f"✅ Session persisted ({len(complete_session)} total messages)\n"
                    )
                else:
                    print(f"⚠️  Session persistence failed\n")

            # STEP 8: Print usage metrics
            if result.usage:
                print(f"📊 Token Usage:")
                print(f"   Input: {result.usage.get('prompt_tokens', 0)}")
                print(f"   Output: {result.usage.get('completion_tokens', 0)}")
                print(f"   Total: {result.usage.get('total_tokens', 0)}\n")

            print("=" * 80)
            print("🏁 AGENT EXECUTION END")
            print("=" * 80 + "\n")

            # Return standardized result
            return {
                "success": result.success,
                "response": result.response,
                "usage": result.usage,
                "model": result.model or input.model_id,
                "finish_reason": result.finish_reason or "stop",
                "run_id": result.run_id,
                "tool_messages": result.tool_messages or [],
                "runtime_type": runtime_type.value,
                "error": result.error,
            }

        except Exception as e:
            print("\n" + "=" * 80)
            print("❌ AGENT EXECUTION FAILED")
            print("=" * 80)
            print(f"Error: {str(e)}")
            print("=" * 80 + "\n")

            logger.error(
                "agent_execution_failed", execution_id=execution_id[:8], error=str(e)
            )

            # Publish critical error as message to the stream
            try:
                error_message = f"❌ Critical Error: {str(e)}"
                turn_number = len(session_history) // 2 + 1 if "session_history" in locals() else 1
                assistant_message_id = f"{execution_id}_assistant_{turn_number}"

                self.control_plane.publish_event(
                    execution_id=execution_id,
                    event_type="message",
                    data={
                        "role": "assistant",
                        "content": error_message,
                        "timestamp": datetime.now(timezone.utc).isoformat(),
                        "message_id": assistant_message_id,
                        "is_error": True,
                    }
                )
                print(f"   📤 Published critical error to stream\n")
            except Exception as publish_error:
                # Don't let error publishing break the error handling
                logger.warning(
                    "failed_to_publish_error_event",
                    execution_id=execution_id[:8],
                    error=str(publish_error)
                )

            return {
                "success": False,
                "error": str(e),
                "model": input.model_id,
                "usage": {},
                "finish_reason": "error",
                "runtime_type": runtime_type.value if "runtime_type" in locals() else "unknown",
            }

    async def _execute_streaming(
        self, runtime, context: RuntimeExecutionContext, input: Any, streaming_helper: StreamingHelper
    ) -> Any:
        """
        Execute with streaming and publish events to Control Plane.

        Args:
            runtime: Runtime instance
            context: Execution context
            input: Original input for additional metadata
            streaming_helper: StreamingHelper instance for tracking events

        Returns:
            Final RuntimeExecutionResult
        """
        # streaming_helper is now passed as parameter instead of created here

        accumulated_response = ""
        final_result = None

        # Define event callback for publishing to Control Plane
        def event_callback(event: Dict):
            """Callback to publish events to Control Plane SSE"""
            event_type = event.get("type")

            if event_type == "content_chunk":
                # Publish content chunk
                streaming_helper.publish_content_chunk(
                    content=event.get("content", ""),
                    message_id=event.get("message_id", context.execution_id),
                )
            elif event_type == "tool_start":
                # Publish tool start
                streaming_helper.publish_tool_start(
                    tool_name=event.get("tool_name"),
                    tool_execution_id=event.get("tool_execution_id"),
                    tool_args=event.get("tool_args", {}),
                    source="agent",
                )
            elif event_type == "tool_complete":
                # Publish tool completion
                streaming_helper.publish_tool_complete(
                    tool_name=event.get("tool_name"),
                    tool_execution_id=event.get("tool_execution_id"),
                    status=event.get("status", "success"),
                    output=event.get("output"),
                    error=event.get("error"),
                    source="agent",
                )

        # Stream execution
        async for chunk in runtime.stream_execute(context, event_callback):
            if chunk.response:
                accumulated_response += chunk.response
                print(chunk.response, end="", flush=True)

            # Keep final result for metadata
            if chunk.usage or chunk.finish_reason:
                final_result = chunk

        print()  # New line after streaming

        # Return final result with accumulated response
        if final_result:
            # Update response with accumulated content
            final_result.response = accumulated_response
            return final_result
        else:
            # Create final result if not provided
            from runtimes.base import RuntimeExecutionResult

            return RuntimeExecutionResult(
                response=accumulated_response,
                usage={},
                success=True,
                finish_reason="stop",
            )

    def _get_framework_name(self, runtime_type: RuntimeType) -> str:
        """
        Get friendly framework name for runtime type.

        Args:
            runtime_type: Runtime type enum

        Returns:
            Framework name string
        """
        mapping = {
            RuntimeType.DEFAULT: "Agno",
            RuntimeType.CLAUDE_CODE: "Claude Code SDK",
        }
        return mapping.get(runtime_type, "Unknown")
