"""
Control Plane Client - Clean API for worker to communicate with Control Plane.

This centralizes all HTTP and WebSocket communication between worker and Control Plane,
providing a clean interface for:
- Event streaming (real-time UI updates via WebSocket or HTTP fallback)
- Session persistence (history storage)
- Metadata caching (execution types)
- Skill resolution
- Bi-directional control messages

Usage:
    from control_plane_client import get_control_plane_client

    client = get_control_plane_client()
    await client.start_websocket()  # If WebSocket enabled
    await client.publish_event_async(execution_id, "message_chunk", {...})
    client.persist_session(execution_id, session_id, user_id, messages)
"""

import os
import httpx
from datetime import datetime, timezone
from typing import Optional, Dict, List, Any
import structlog

logger = structlog.get_logger()


class ControlPlaneClient:
    """Client for communicating with the Control Plane API from workers."""

    def __init__(
        self,
        base_url: str,
        api_key: str,
        websocket_enabled: bool = False,
        websocket_url: Optional[str] = None,
        worker_id: Optional[str] = None
    ):
        """
        Initialize Control Plane client.

        Args:
            base_url: Control Plane URL (e.g., http://localhost:8000)
            api_key: Kubiya API key for authentication
            websocket_enabled: Whether WebSocket is enabled
            websocket_url: WebSocket URL if enabled
            worker_id: Worker ID for WebSocket connection
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.headers = {"Authorization": f"UserKey {api_key}"}

        # Use BOTH sync and async clients for different use cases
        # Sync client for backwards compatibility with non-async code
        self._client = httpx.Client(
            timeout=httpx.Timeout(30.0, connect=5.0, read=30.0, write=10.0),
            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
        )

        # Async client for streaming/real-time operations
        # Longer read timeout to handle streaming scenarios
        self._async_client = httpx.AsyncClient(
            timeout=httpx.Timeout(60.0, connect=5.0, read=60.0, write=10.0),
            limits=httpx.Limits(max_connections=20, max_keepalive_connections=10),
        )

        # WebSocket client for persistent connection
        self.websocket_client: Optional[Any] = None

        # Initialize WebSocket client if enabled and environment supports it
        if websocket_enabled and websocket_url and worker_id:
            from control_plane_api.worker.utils.environment import should_use_websocket

            if should_use_websocket():
                from control_plane_api.worker.websocket_client import WorkerWebSocketClient

                self.websocket_client = WorkerWebSocketClient(
                    worker_id=worker_id,
                    websocket_url=websocket_url,
                    api_key=api_key,
                    on_control_message=self._handle_control_message
                )
                logger.info("websocket_client_initialized", worker_id=worker_id[:8])
            else:
                logger.info("websocket_skipped_serverless_environment")

    def __del__(self):
        """Close the HTTP clients on cleanup."""
        try:
            self._client.close()
        except:
            pass
        # Async client cleanup happens via context manager or explicit close

    async def aclose(self):
        """Async cleanup for async client."""
        try:
            await self._async_client.aclose()
        except:
            pass

    async def start_websocket(self):
        """Start WebSocket client if enabled."""
        if self.websocket_client:
            await self.websocket_client.start()
            logger.info("websocket_started")

    async def stop_websocket(self):
        """Stop WebSocket client if running."""
        if self.websocket_client:
            await self.websocket_client.stop()
            logger.info("websocket_stopped")

    def publish_event(
        self,
        execution_id: str,
        event_type: str,
        data: Dict[str, Any],
    ) -> bool:
        """
        Publish a streaming event for real-time UI updates (SYNC version).

        NOTE: This is the BLOCKING version. For real-time streaming,
        use publish_event_async() instead to avoid blocking the event loop.

        Args:
            execution_id: Execution ID
            event_type: Event type (message_chunk, tool_started, etc.)
            data: Event payload

        Returns:
            True if successful, False otherwise
        """
        try:
            # Sanitize data to remove non-JSON-serializable objects
            import json
            import asyncio

            def sanitize_value(val):
                """Remove non-JSON-serializable objects"""
                try:
                    # Fast path for JSON primitives
                    if val is None or isinstance(val, (bool, int, float, str)):
                        return val

                    # Check type name to avoid event loop issues
                    type_name = type(val).__name__
                    type_module = str(type(val).__module__)

                    # Remove asyncio objects by checking module and type name
                    if 'asyncio' in type_module or any(x in type_name for x in ['Event', 'Lock', 'Queue', 'Semaphore', 'Condition']):
                        return f"<{type_name}>"
                    elif isinstance(val, dict):
                        return {k: sanitize_value(v) for k, v in val.items()}
                    elif isinstance(val, (list, tuple)):
                        return [sanitize_value(v) for v in val]
                    else:
                        try:
                            json.dumps(val)
                            return val
                        except (TypeError, ValueError, RuntimeError):
                            # RuntimeError catches "bound to different event loop" errors
                            return f"<non-serializable: {type_name}>"
                except Exception as e:
                    # Catch-all for ANY errors during sanitization itself
                    # Do NOT attempt to inspect the value here - it may cause event loop errors
                    return "<sanitization-error>"

            sanitized_data = sanitize_value(data)

            url = f"{self.base_url}/api/v1/executions/{execution_id}/events"
            payload = {
                "event_type": event_type,
                "data": sanitized_data,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

            # Double-check: Try to serialize the payload before sending
            try:
                json.dumps(payload)
            except Exception as serialize_err:
                logger.error(
                    "payload_serialization_test_failed",
                    execution_id=execution_id[:8],
                    event_type=event_type,
                    error=str(serialize_err)[:200],
                )
                # If we can't serialize it, don't even try to send
                return False

            response = self._client.post(url, json=payload, headers=self.headers)

            if response.status_code not in (200, 202):
                logger.warning(
                    "event_publish_failed",
                    status=response.status_code,
                    execution_id=execution_id[:8],
                    event_type=event_type,
                )
                return False

            return True

        except Exception as e:
            # Sanitize error message to avoid serialization issues
            import re
            error_str = str(e)
            # Remove asyncio object references that cause serialization errors
            error_str = re.sub(r'<asyncio\.\w+\.\w+ object at 0x[0-9a-f]+ \[[\w\s]+\]>', '[asyncio-object]', error_str)

            logger.warning(
                "event_publish_error",
                error=error_str[:500],  # Truncate to prevent huge error messages
                execution_id=execution_id[:8],
                event_type=event_type,
            )
            return False

    async def publish_event_async(
        self,
        execution_id: str,
        event_type: str,
        data: Dict[str, Any],
    ) -> bool:
        """
        Publish a streaming event for real-time UI updates (ASYNC version).

        This is NON-BLOCKING and uses WebSocket-first strategy:
        1. Try WebSocket if connected
        2. Fallback to HTTP if WebSocket unavailable or queue full

        Args:
            execution_id: Execution ID
            event_type: Event type (message_chunk, tool_started, etc.)
            data: Event payload

        Returns:
            True if successful, False otherwise
        """
        # Try WebSocket first if available and connected
        if self.websocket_client and self.websocket_client.is_connected():
            success = await self.websocket_client.send_event(execution_id, event_type, data)
            if success:
                return True

            # Queue full - fallback to HTTP immediately
            logger.warning("websocket_queue_full_fallback_http", execution_id=execution_id[:8])

        # Fallback to HTTP
        return await self._publish_event_http(execution_id, event_type, data)

    async def _publish_event_http(
        self,
        execution_id: str,
        event_type: str,
        data: Dict[str, Any],
    ) -> bool:
        """
        Publish event via HTTP (internal method for fallback).

        Args:
            execution_id: Execution ID
            event_type: Event type
            data: Event payload

        Returns:
            True if successful, False otherwise
        """
        try:
            # Sanitize data to remove non-JSON-serializable objects
            import json
            import asyncio

            def sanitize_value(val):
                """Remove non-JSON-serializable objects"""
                try:
                    # Fast path for JSON primitives
                    if val is None or isinstance(val, (bool, int, float, str)):
                        return val

                    # Check type name to avoid event loop issues
                    type_name = type(val).__name__
                    type_module = str(type(val).__module__)

                    # Remove asyncio objects by checking module and type name
                    if 'asyncio' in type_module or any(x in type_name for x in ['Event', 'Lock', 'Queue', 'Semaphore', 'Condition']):
                        return f"<{type_name}>"
                    elif isinstance(val, dict):
                        return {k: sanitize_value(v) for k, v in val.items()}
                    elif isinstance(val, (list, tuple)):
                        return [sanitize_value(v) for v in val]
                    else:
                        try:
                            json.dumps(val)
                            return val
                        except (TypeError, ValueError, RuntimeError):
                            # RuntimeError catches "bound to different event loop" errors
                            return f"<non-serializable: {type_name}>"
                except Exception as e:
                    # Catch-all for ANY errors during sanitization itself
                    # Do NOT attempt to inspect the value here - it may cause event loop errors
                    return "<sanitization-error>"

            sanitized_data = sanitize_value(data)

            url = f"{self.base_url}/api/v1/executions/{execution_id}/events"
            payload = {
                "event_type": event_type,
                "data": sanitized_data,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

            # Double-check: Try to serialize the payload before sending
            try:
                json.dumps(payload)
            except Exception as serialize_err:
                logger.error(
                    "payload_serialization_test_failed",
                    execution_id=execution_id[:8],
                    event_type=event_type,
                    error=str(serialize_err)[:200],
                )
                # If we can't serialize it, don't even try to send
                return False

            response = await self._async_client.post(url, json=payload, headers=self.headers)

            if response.status_code not in (200, 202):
                logger.warning(
                    "event_publish_failed",
                    status=response.status_code,
                    execution_id=execution_id[:8],
                    event_type=event_type,
                )
                return False

            return True

        except Exception as e:
            # Sanitize error message to avoid serialization issues
            import re
            error_str = str(e)
            # Remove asyncio object references that cause serialization errors
            error_str = re.sub(r'<asyncio\.\w+\.\w+ object at 0x[0-9a-f]+ \[[\w\s]+\]>', '[asyncio-object]', error_str)

            logger.warning(
                "event_publish_error",
                error=error_str[:500],  # Truncate to prevent huge error messages
                execution_id=execution_id[:8],
                event_type=event_type,
            )
            return False

    def cache_metadata(
        self,
        execution_id: str,
        execution_type: str,
    ) -> bool:
        """
        Cache execution metadata in Redis for fast SSE lookups.

        This eliminates the need for database queries on every SSE connection.

        Args:
            execution_id: Execution ID
            execution_type: "AGENT" or "TEAM"

        Returns:
            True if successful, False otherwise
        """
        return self.publish_event(
            execution_id=execution_id,
            event_type="metadata",
            data={"execution_type": execution_type},
        )

    def get_session(
        self,
        execution_id: str,
        session_id: Optional[str] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Retrieve session history from Control Plane database.

        This loads conversation history so workers can restore context
        across multiple execution turns.

        Args:
            execution_id: Execution ID
            session_id: Session ID (defaults to execution_id if not provided)

        Returns:
            Dict with session data including messages, or None if not found
        """
        try:
            session_id = session_id or execution_id
            url = f"{self.base_url}/api/v1/executions/{execution_id}/session"

            response = self._client.get(url, headers=self.headers)

            if response.status_code == 200:
                session_data = response.json()
                logger.info(
                    "session_loaded",
                    execution_id=execution_id[:8],
                    message_count=len(session_data.get("messages", [])),
                )
                return session_data
            elif response.status_code == 404:
                logger.info(
                    "session_not_found",
                    execution_id=execution_id[:8],
                )
                return None
            else:
                logger.warning(
                    "session_load_failed",
                    status=response.status_code,
                    execution_id=execution_id[:8],
                )
                return None

        except Exception as e:
            logger.warning(
                "session_load_error",
                error=str(e),
                execution_id=execution_id[:8],
            )
            return None

    def persist_session(
        self,
        execution_id: str,
        session_id: str,
        user_id: Optional[str],
        messages: List[Dict[str, Any]],
        metadata: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """
        Persist session history to Control Plane database.

        This ensures history is available even when worker is offline.

        Args:
            execution_id: Execution ID
            session_id: Session ID
            user_id: User ID
            messages: List of session messages
            metadata: Optional metadata

        Returns:
            True if successful, False otherwise
        """
        try:
            url = f"{self.base_url}/api/v1/executions/{execution_id}/session"
            payload = {
                "session_id": session_id,
                "user_id": user_id,
                "messages": messages,
                "metadata": metadata or {},
            }

            response = self._client.post(url, json=payload, headers=self.headers)

            if response.status_code in (200, 201):
                logger.info(
                    "session_persisted",
                    execution_id=execution_id[:8],
                    message_count=len(messages),
                )
                return True
            else:
                logger.warning(
                    "session_persistence_failed",
                    status=response.status_code,
                    execution_id=execution_id[:8],
                )
                return False

        except Exception as e:
            logger.warning(
                "session_persistence_error",
                error=str(e),
                execution_id=execution_id[:8],
            )
            return False

    def get_skills(
        self,
        agent_id: str,
    ) -> List[Dict[str, Any]]:
        """
        Fetch resolved skills for an agent from Control Plane.

        This endpoint returns skills merged from all layers:
        - All agent environments (many-to-many)
        - Team skills (if agent has team)
        - All team environments (many-to-many)
        - Agent's own skills

        Args:
            agent_id: Agent ID

        Returns:
            List of skill configurations with source and inheritance info
        """
        try:
            url = f"{self.base_url}/api/v1/skills/associations/agents/{agent_id}/skills/resolved"
            response = self._client.get(url, headers=self.headers)

            if response.status_code == 200:
                skills = response.json()
                logger.info(
                    "skills_fetched",
                    agent_id=agent_id[:8],
                    skill_count=len(skills),
                )
                return skills
            else:
                logger.warning(
                    "skills_fetch_failed",
                    status=response.status_code,
                    agent_id=agent_id[:8],
                )
                return []

        except Exception as e:
            logger.warning(
                "skills_fetch_error",
                error=str(e),
                agent_id=agent_id[:8],
            )
            return []

    def get_team_skills(
        self,
        team_id: str,
    ) -> List[Dict[str, Any]]:
        """
        Fetch resolved skills for a team from Control Plane.

        This endpoint returns skills merged from all layers:
        - All team environments (many-to-many)
        - Team's own skills

        Args:
            team_id: Team ID

        Returns:
            List of skill configurations with source and inheritance info
        """
        try:
            url = f"{self.base_url}/api/v1/skills/associations/teams/{team_id}/skills/resolved"
            response = self._client.get(url, headers=self.headers)

            if response.status_code == 200:
                skills = response.json()
                logger.info(
                    "team_skills_fetched",
                    team_id=team_id[:8],
                    skill_count=len(skills),
                )
                return skills
            else:
                logger.warning(
                    "team_skills_fetch_failed",
                    status=response.status_code,
                    team_id=team_id[:8],
                )
                return []

        except Exception as e:
            logger.warning(
                "team_skills_fetch_error",
                error=str(e),
                team_id=team_id[:8],
            )
            return []

    def get_agent_execution_environment(
        self,
        agent_id: str,
    ) -> Dict[str, str]:
        """
        Fetch resolved execution environment for an agent from Control Plane.

        This endpoint returns a fully resolved environment variable dict with:
        - Custom env vars from agent configuration
        - Secret values (resolved from Kubiya vault)
        - Integration tokens (resolved and mapped to env var names like GH_TOKEN, JIRA_TOKEN)

        Args:
            agent_id: Agent ID

        Returns:
            Dict of environment variables ready to inject into agent execution
        """
        try:
            url = f"{self.base_url}/api/v1/execution-environment/agents/{agent_id}/resolved"
            response = self._client.get(url, headers=self.headers)

            if response.status_code == 200:
                env_vars = response.json()
                logger.info(
                    "agent_execution_environment_fetched",
                    agent_id=agent_id[:8],
                    env_var_count=len(env_vars),
                    env_var_keys=list(env_vars.keys()),
                )
                return env_vars
            else:
                logger.warning(
                    "agent_execution_environment_fetch_failed",
                    status=response.status_code,
                    agent_id=agent_id[:8],
                )
                return {}

        except Exception as e:
            logger.warning(
                "agent_execution_environment_fetch_error",
                error=str(e),
                agent_id=agent_id[:8],
            )
            return {}

    def get_team_execution_environment(
        self,
        team_id: str,
    ) -> Dict[str, str]:
        """
        Fetch resolved execution environment for a team from Control Plane.

        This endpoint returns a fully resolved environment variable dict with:
        - Custom env vars from team configuration
        - Secret values (resolved from Kubiya vault)
        - Integration tokens (resolved and mapped to env var names like GH_TOKEN, JIRA_TOKEN)

        Args:
            team_id: Team ID

        Returns:
            Dict of environment variables ready to inject into team execution
        """
        try:
            url = f"{self.base_url}/api/v1/execution-environment/teams/{team_id}/resolved"
            response = self._client.get(url, headers=self.headers)

            if response.status_code == 200:
                env_vars = response.json()
                logger.info(
                    "team_execution_environment_fetched",
                    team_id=team_id[:8],
                    env_var_count=len(env_vars),
                    env_var_keys=list(env_vars.keys()),
                )
                return env_vars
            else:
                logger.warning(
                    "team_execution_environment_fetch_failed",
                    status=response.status_code,
                    team_id=team_id[:8],
                )
                return {}

        except Exception as e:
            logger.warning(
                "team_execution_environment_fetch_error",
                error=str(e),
                team_id=team_id[:8],
            )
            return {}

    async def create_job_execution_record(
        self,
        execution_id: str,
        job_id: Optional[str],
        organization_id: str,
        entity_type: str,
        entity_id: Optional[str],
        prompt: str,
        trigger_type: str,
        trigger_metadata: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        Create execution and job_executions records for a scheduled job.

        This calls the Control Plane API to create execution records
        instead of directly accessing Supabase.

        Args:
            execution_id: Execution ID
            job_id: Job ID (optional)
            organization_id: Organization ID
            entity_type: "agent" or "team"
            entity_id: Agent or team ID
            prompt: Prompt text
            trigger_type: "cron", "webhook", or "manual"
            trigger_metadata: Additional trigger metadata

        Returns:
            Dict with execution_id, status, and created_at
        """
        try:
            url = f"{self.base_url}/api/v1/executions/create"
            payload = {
                "execution_id": execution_id,
                "job_id": job_id,
                "organization_id": organization_id,
                "entity_type": entity_type,
                "entity_id": entity_id,
                "prompt": prompt,
                "trigger_type": trigger_type,
                "trigger_metadata": trigger_metadata,
            }

            response = await self._async_client.post(url, json=payload, headers=self.headers)

            if response.status_code == 201:
                result = response.json()
                logger.info(
                    "job_execution_record_created",
                    execution_id=execution_id[:8],
                    job_id=job_id[:8] if job_id else None,
                )
                return result
            else:
                logger.error(
                    "job_execution_record_creation_failed",
                    status=response.status_code,
                    execution_id=execution_id[:8],
                    response=response.text,
                )
                raise Exception(f"Failed to create execution record: HTTP {response.status_code}")

        except Exception as e:
            logger.error(
                "job_execution_record_creation_error",
                error=str(e),
                execution_id=execution_id[:8],
            )
            raise

    async def update_job_execution_status(
        self,
        execution_id: str,
        job_id: str,
        status: str,
        duration_ms: Optional[int] = None,
        error_message: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Update job_executions record with execution results.

        This calls the Control Plane API to update job execution status
        instead of directly accessing Supabase.

        Args:
            execution_id: Execution ID
            job_id: Job ID
            status: Final status ("completed" or "failed")
            duration_ms: Execution duration in milliseconds
            error_message: Error message if failed

        Returns:
            Dict with job_id, execution_id, and status
        """
        try:
            url = f"{self.base_url}/api/v1/executions/{execution_id}/job/{job_id}/status"
            payload = {
                "status": status,
                "duration_ms": duration_ms,
                "error_message": error_message,
            }

            response = await self._async_client.post(url, json=payload, headers=self.headers)

            if response.status_code == 200:
                result = response.json()
                logger.info(
                    "job_execution_status_updated",
                    execution_id=execution_id[:8],
                    job_id=job_id[:8],
                    status=status,
                )
                return result
            else:
                logger.error(
                    "job_execution_status_update_failed",
                    status_code=response.status_code,
                    execution_id=execution_id[:8],
                    job_id=job_id[:8],
                    response=response.text,
                )
                raise Exception(f"Failed to update job execution status: HTTP {response.status_code}")

        except Exception as e:
            logger.error(
                "job_execution_status_update_error",
                error=str(e),
                execution_id=execution_id[:8],
                job_id=job_id[:8],
            )
            raise

    async def _handle_control_message(self, message: Dict[str, Any]):
        """
        Handle control messages from control plane via WebSocket.

        This method is called when the worker receives a control message
        from the control plane (pause, resume, cancel, reload_config).

        Args:
            message: Control message with command, execution_id, and data
        """
        command = message.get("command")
        execution_id = message.get("execution_id")

        try:
            # Import Temporal client here to avoid circular import
            from control_plane_api.app.lib.temporal_client import get_temporal_client

            temporal_client = get_temporal_client()
            workflow_handle = temporal_client.get_workflow_handle(execution_id)

            if command == "pause":
                await workflow_handle.signal("pause_execution")
                logger.info("control_command_executed", command="pause", execution_id=execution_id[:8])

            elif command == "resume":
                await workflow_handle.signal("resume_execution")
                logger.info("control_command_executed", command="resume", execution_id=execution_id[:8])

            elif command == "cancel":
                await workflow_handle.cancel()
                logger.info("control_command_executed", command="cancel", execution_id=execution_id[:8])

            elif command == "reload_config":
                # Future: Reload config without restart
                logger.info("control_command_not_implemented", command="reload_config", execution_id=execution_id[:8])

            else:
                logger.warning("unknown_control_command", command=command, execution_id=execution_id[:8])

        except Exception as e:
            logger.error(
                "control_command_error",
                error=str(e),
                command=command,
                execution_id=execution_id[:8] if execution_id else None
            )


# Singleton instance
_control_plane_client: Optional[ControlPlaneClient] = None


def get_control_plane_client() -> ControlPlaneClient:
    """
    Get or create the Control Plane client singleton.

    Reads configuration from environment variables:
    - CONTROL_PLANE_URL: Control Plane URL
    - KUBIYA_API_KEY: API key for authentication
    - WEBSOCKET_ENABLED: Whether WebSocket is enabled (from registration)
    - WEBSOCKET_URL: WebSocket URL (from registration)
    - WORKER_ID: Worker ID (from registration)

    Returns:
        ControlPlaneClient instance

    Raises:
        ValueError: If required environment variables are not set
    """
    global _control_plane_client

    if _control_plane_client is None:
        base_url = os.environ.get("CONTROL_PLANE_URL")
        api_key = os.environ.get("KUBIYA_API_KEY")

        # WebSocket config from environment (set by worker.py after registration)
        websocket_enabled = os.environ.get("WEBSOCKET_ENABLED", "false").lower() == "true"
        websocket_url = os.environ.get("WEBSOCKET_URL")
        worker_id = os.environ.get("WORKER_ID")

        if not base_url:
            raise ValueError("CONTROL_PLANE_URL environment variable not set")
        if not api_key:
            raise ValueError("KUBIYA_API_KEY environment variable not set")

        _control_plane_client = ControlPlaneClient(
            base_url=base_url,
            api_key=api_key,
            websocket_enabled=websocket_enabled,
            websocket_url=websocket_url,
            worker_id=worker_id
        )

        logger.info(
            "control_plane_client_initialized",
            base_url=base_url,
            websocket_enabled=websocket_enabled,
        )

    return _control_plane_client
