import signal
import os
import time
import json
import asyncio
from pathlib import Path
from blocks_control_sdk.utils.logger import log
from enum import Enum
from types import FunctionType
from typing import Any, Tuple, Union, List, Optional, Dict, TYPE_CHECKING, AsyncIterator
from blocks.utils import BackgroundCommandOutput
import uuid
from pydantic import BaseModel, ConfigDict

if TYPE_CHECKING:
    from litellm.types.utils import Message, ModelResponse

class NotificationsV2:
    NOTIFY_COMPLETE_V2 = "notify_complete_v2"
    NOTIFY_MESSAGE_V2 = "notify_message_v2"
    NOTIFY_TOOL_CALL_V2 = "notify_tool_call_v2"
    NOTIFY_START_V2 = "notify_start_v2"
    NOTIFY_RESUME_V2 = "notify_resume_v2"
    NOTIFY_CHAT_THREAD_ID_UPDATE = "notify_chat_thread_id_update"

class NotifyBase(BaseModel):
    type: str
    status: Optional[str] = None
    chat_thread_id: Optional[str] = None

class NotifyMessageArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_MESSAGE_V2
    message: Any  # Accepts both local Message and litellm Message

class NotifyCompleteArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_COMPLETE_V2
    last_message: str

class NotifyToolCallArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_TOOL_CALL_V2
    tool_name: str
    serialized_args: str

class NotifyStartArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_START_V2

class NotifyResumeArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_RESUME_V2

class NotifyChatThreadIdUpdateArgs(NotifyBase):
    type: str = NotificationsV2.NOTIFY_CHAT_THREAD_ID_UPDATE
    new_chat_thread_id: str
    old_chat_thread_id: Optional[str] = None

class LLM(Enum):
    CLAUDE = "claude"
    CODEX = "codex"
    GEMINI = "gemini"

LLM_API_KEYS = {
    LLM.CLAUDE: os.getenv("ANTHROPIC_API_KEY"),
    LLM.CODEX: os.getenv("OPENAI_API_KEY"),
    LLM.GEMINI: os.getenv("GEMINI_API_KEY"),
}

class CodingAgentBaseCLI():

    class AgentStatus:
        TURNS_COMPLETED = "turns_completed"
        TURNS_IN_PROGRESS = "turns_in_progress"
        TURNS_NOT_STARTED = "turns_not_started"

    class Notifications:
        NOTIFY_START_V2 = "notify_start_v2"
        NOTIFY_RESUME_V2 = "notify_resume_v2"
        NOTIFY_COMPLETE_V2 = "notify_complete_v2"
        NOTIFY_MESSAGE_V2 = "notify_message_v2"
        NOTIFY_TODO_LIST_UPDATE = "notify_todo_list_update"
        NOTIFY_TOOL_CALL_V2 = "notify_tool_call_v2"
        NOTIFY_CHAT_THREAD_ID_UPDATE = "notify_chat_thread_id_update"
    
    pid: int = None
    is_session_active: bool = False
    status: str = AgentStatus.TURNS_NOT_STARTED
    is_interrupt_triggered: bool = False
    chat_thread_id: str = None
    workspace_dir: Optional[Path] = None

    def __init__(self, chat_thread_id: str = None, workspace_dir: str = None):
        self.notifications = self.Notifications()
        self._notification_callbacks = {}
        self.assistant_messages = []  # Initialize as instance variable, not class variable
        self.callback_errors = []  # Track any errors that occur in notification callbacks

        # Raw message capture for debugging and regression testing
        self.raw_messages: List[Dict[str, Any]] = []
        self._raw_message_index: int = 0

        self.set_chat_thread_id(chat_thread_id)

        # if workspace_dir is None set to cwd
        if workspace_dir is not None:
            self.workspace_dir = Path(workspace_dir)
        else:
            self.workspace_dir = Path.cwd()

    def tool_call_arg_kb_size(self, content: str) -> int:
        content = content or ""
        size_bytes = len(content.encode('utf-8'))
        size_kb = size_bytes / 1024
        return int(size_kb)

    def tool_call_arg_kb_size_truncate_to_limit(self, content: str, kb_limit: int) -> str:
        limit_bytes = kb_limit * 1024
        encoded = content.encode("utf-8")
        if len(encoded) <= limit_bytes:
            return content
        # cut the bytes safely
        truncated_bytes = encoded[:limit_bytes]
        # decode with ignore to drop partial utf-8 chars
        truncated_str = truncated_bytes.decode("utf-8", errors="ignore")
        return truncated_str + "..."

    @staticmethod
    def resolve_agent_class(provider: LLM):
        from .agent_codex import Codex
        from .agent_claude_exp import ClaudeCode
        from .agent_gemini_exp import GeminiCLI

        if provider == LLM.CODEX:
            return Codex
        elif provider == LLM.CLAUDE:
            return ClaudeCode
        elif provider == LLM.GEMINI:
            return GeminiCLI
        else:
            raise ValueError(f"Unknown provider: {provider}")

    @staticmethod
    def resolve_provider_from_keys(default_provider: Union[LLM, None] = None) -> LLM:

        if default_provider is not None and default_provider in LLM_API_KEYS:
            print(f"Attempting to resolve LLM provider from default_provider in automation: {default_provider}")
            if LLM_API_KEYS.get(default_provider) is not None:
                return default_provider
            else:
                raise ValueError(
                    f"LLM API key for {default_provider} is not set, you must set the environment variable {default_provider.value}"
                )

        print("Attempting to resolve LLM provider from api keys...")

        if LLM_API_KEYS.get(LLM.CLAUDE) is not None:
            return LLM.CLAUDE
        elif LLM_API_KEYS.get(LLM.CODEX) is not None:
            return LLM.CODEX
        elif LLM_API_KEYS.get(LLM.GEMINI) is not None:
            return LLM.GEMINI
        else:
            llm_name_strs = [llm.value for llm in LLM]
            raise ValueError(
                "No LLM API keys set, you must set at least one of the following environment variables: "
                f"{', '.join(llm_name_strs)}"
            )

    def _kill_process(self, pid: int):
        try:
            os.killpg(pid, signal.SIGTERM)
            time.sleep(1) 
            os.killpg(pid, signal.SIGKILL) 
        except ProcessLookupError:
            pass
        except PermissionError:
            # Fallback to killing just the main process
            try:
                os.kill(pid, signal.SIGKILL)
            except ProcessLookupError:
                pass

    def register_notification(self, notification_type: str, callback):
        """Register a callback for a specific notification type"""
        if notification_type not in self._notification_callbacks:
            self._notification_callbacks[notification_type] = []
        self._notification_callbacks[notification_type].append(callback)

    def notify(self, arg: NotifyBase):
        """Trigger all callbacks registered for a notification type with args and kwargs"""

        notification_type = arg.type

        arg.chat_thread_id = self.chat_thread_id
        arg.status = self.status

        size = notification_type in self._notification_callbacks and len(self._notification_callbacks[notification_type]) or 0

        log.debug(f">Notifying callbacks for notification type: {notification_type}, #callbacks: {size}")

        if notification_type in self._notification_callbacks:
            for callback in self._notification_callbacks[notification_type]:
                if callable(callback):
                    try:
                        callback(arg)
                    except Exception as e:
                        error_info = {
                            "notification_type": notification_type,
                            "callback": str(callback),
                            "error": str(e)
                        }
                        self.callback_errors.append(error_info)
                        log.error(f">Error calling callback {callback} for notification type {notification_type}: {e}")
                else:
                    log.error(f">Callback {callback} for notification type {notification_type} is not callable")

    def clear_messages(self):
        self.assistant_messages = []
        self.clear_raw_messages()

    def clear_raw_messages(self):
        """Clear raw message capture buffer."""
        self.raw_messages = []
        self._raw_message_index = 0

    def capture_raw_line(
        self,
        raw_json: Dict[str, Any],
        parsed_type: str = "unknown",
        parse_success: bool = True,
        parse_error: Optional[str] = None
    ):
        """
        Capture a raw JSON line for debugging and regression testing.

        Args:
            raw_json: The raw JSON object parsed from the agent's output
            parsed_type: Type of parsed content ("message", "tool_call", "parse_error", "other")
            parse_success: Whether parsing was successful
            parse_error: Error message if parsing failed
        """
        entry = {
            "index": self._raw_message_index,
            "timestamp": time.time(),
            "raw_json": raw_json,
            "parsed_type": parsed_type,
            "parse_success": parse_success,
            "parse_error": parse_error,
        }

        self.raw_messages.append(entry)
        self._raw_message_index += 1

    def set_chat_thread_id(self, chat_thread_id: str = None):
        new_chat_thread_id = chat_thread_id or str(uuid.uuid4())
        old_chat_thread_id = self.chat_thread_id
        self.notify(NotifyChatThreadIdUpdateArgs(
            new_chat_thread_id=new_chat_thread_id,
            old_chat_thread_id=old_chat_thread_id
        ))
        self.chat_thread_id = new_chat_thread_id

    def new_chat_thread(self, chat_thread_id: str = None):
        self.set_chat_thread_id(chat_thread_id)
        self.clear_messages()
    
    def notify_complete_callback(self, text: str):
        """Default callback method that can be overridden"""
        pass

    def _start(self):
        self.is_session_active = True

    def interrupt(self):
        if self.pid is not None:
            self._kill_process(self.pid)
            self.pid = None
            self.is_interrupt_triggered = True

    def query(self, query: str) -> BackgroundCommandOutput:
        pass

    def query_sync(self, prompt: str) -> str:
        """
        Synchronous query that runs the agent and returns the last message.

        Args:
            prompt: The query/prompt to send to the agent

        Returns:
            The content of the last assistant message
        """
        last_message: Optional[str] = None
        is_complete = False

        def on_message(notification: NotifyMessageArgs):
            nonlocal last_message
            message = notification.message
            last_message = message

        def on_complete(notification: NotifyCompleteArgs):
            nonlocal is_complete
            is_complete = True

        # Register callbacks
        self.register_notification(self.Notifications.NOTIFY_MESSAGE_V2, on_message)
        self.register_notification(self.Notifications.NOTIFY_COMPLETE_V2, on_complete)

        try:
            # Run the query and wait for completion
            self.query(prompt)
            while not is_complete:
                time.sleep(0.2)
            
        finally:
            # Clean up callbacks
            if self.Notifications.NOTIFY_MESSAGE_V2 in self._notification_callbacks:
                try:
                    self._notification_callbacks[self.Notifications.NOTIFY_MESSAGE_V2].remove(on_message)
                except ValueError:
                    pass
            if self.Notifications.NOTIFY_COMPLETE_V2 in self._notification_callbacks:
                try:
                    self._notification_callbacks[self.Notifications.NOTIFY_COMPLETE_V2].remove(on_complete)
                except ValueError:
                    pass

        return last_message

    async def stream(self, prompt: str) -> AsyncIterator["Message"]:
        """
        Async generator that yields Message objects as they stream in.

        Usage:
            async for message in agent.stream("Tell me a joke"):
                print(message.content)

        Args:
            prompt: The query/prompt to send to the agent

        Yields:
            Message objects containing assistant messages
        """
        queue: asyncio.Queue = asyncio.Queue()
        loop = asyncio.get_event_loop()

        def on_message(notification: NotifyMessageArgs):
            """Bridge from sync callback to async queue"""
            loop.call_soon_threadsafe(queue.put_nowait, ("message", notification.message))

        # on tool call
        def on_tool_call(notification: NotifyToolCallArgs):
            """Bridge from sync callback to async queue"""
            loop.call_soon_threadsafe(queue.put_nowait, ("tool_call", (notification.tool_name, notification.serialized_args)))

        def on_complete(notification: NotifyCompleteArgs):
            """Signal completion by pushing sentinel"""
            loop.call_soon_threadsafe(queue.put_nowait, ("complete", notification.last_message))

        # Register callbacks for streaming
        self.register_notification(self.Notifications.NOTIFY_MESSAGE_V2, on_message)
        self.register_notification(self.Notifications.NOTIFY_COMPLETE_V2, on_complete)
        self.register_notification(self.Notifications.NOTIFY_TOOL_CALL_V2, on_tool_call)
        
        try:
            # Start the query (non-blocking, returns immediately)
            self.query(prompt)

            # Yield messages until completion
            while True:
                event_type, payload = await queue.get()
                if event_type == "complete":
                    break
                elif event_type == "tool_call":
                    yield payload
                elif event_type == "message":
                    yield payload
        finally:
            # Clean up callbacks to prevent memory leaks on repeated calls
            if self.Notifications.NOTIFY_MESSAGE_V2 in self._notification_callbacks:
                try:
                    self._notification_callbacks[self.Notifications.NOTIFY_MESSAGE_V2].remove(on_message)
                except ValueError:
                    pass
            if self.Notifications.NOTIFY_COMPLETE_V2 in self._notification_callbacks:
                try:
                    self._notification_callbacks[self.Notifications.NOTIFY_COMPLETE_V2].remove(on_complete)
                except ValueError:
                    pass
            if self.Notifications.NOTIFY_TOOL_CALL_V2 in self._notification_callbacks:
                try:
                    self._notification_callbacks[self.Notifications.NOTIFY_TOOL_CALL_V2].remove(on_tool_call)
                except ValueError:
                    pass

    def warm_up_mcp(self, default_packages: List[str] = None) -> Dict[str, Any]:
        """
        Pre-install MCP (Model Context Protocol) packages that use npx.
        Reads configuration from ~/.config/blocks/mcp.json and installs
        all npx-based packages globally using npm.
        
        Args:
            default_packages: List of default packages to always install (with optional versions)
                             e.g., ["mcp-remote"]
        
        Returns:
            Dict containing status, installed packages, and any errors
        """
        from blocks import bash
        
        # Default package(s) to always install
        if default_packages is None:
            default_packages = ["mcp-remote"]
        
        result = {
            "status": "success",
            "packages": [],
            "errors": [],
            "message": ""
        }
        
        try:
            # Construct path to MCP config
            mcp_config_path = Path.home() / ".config" / "blocks" / "mcp.json"
            
            # Check if config file exists
            if not mcp_config_path.exists():
                result["status"] = "skipped"
                result["message"] = f"MCP config not found at {mcp_config_path}"
                print(f"[warm_up_mcp] {result['message']}")
                return result
            
            # Read and parse MCP configuration
            with open(mcp_config_path, "r") as f:
                mcp_config = json.load(f)
            
            if not mcp_config:
                result["status"] = "skipped"
                result["message"] = "MCP config is empty"
                print(f"[warm_up_mcp] {result['message']}")
                return result
            
            # Extract npx-based packages
            npx_packages = []
            for server_name, server_config in mcp_config.items():
                # Skip disabled servers
                if server_config.get("disabled", False):
                    print(f"[warm_up_mcp] Skipping disabled server: {server_name}")
                    continue
                
                # Check if this is an npx-based server
                if server_config.get("command") == "npx":
                    args = server_config.get("args", [])
                    
                    # Extract package name from args
                    # Typical format: ["-y", "@package/name@version"]
                    package_name = None
                    for arg in args:
                        # Skip flags
                        if arg.startswith("-"):
                            continue
                        # This should be the package name
                        package_name = arg
                        break
                    
                    if package_name:
                        # Keep the package name as-is, including version if specified
                        # Examples:
                        # - "@modelcontextprotocol/server-slack" (no version)
                        # - "@modelcontextprotocol/server-slack@latest" (with version)
                        # - "firecrawl-mcp@1.2.3" (non-scoped with version)
                        
                        npx_packages.append({
                            "name": package_name,
                            "server": server_name
                        })
                        print(f"[warm_up_mcp] Found npx package: {package_name} (server: {server_name})")
            
            # Add default packages
            for default_pkg in default_packages:
                print(f"[warm_up_mcp] Adding default package: {default_pkg}")
            
            # Prepare npm install command
            package_names = [pkg["name"] for pkg in npx_packages] + default_packages
            
            if not package_names:
                result["status"] = "skipped"
                result["message"] = "No packages to install"
                print(f"[warm_up_mcp] {result['message']}")
                return result
            
            result["packages"] = package_names
            
            # Run npm install globally
            install_command = f"npm install -g {' '.join(package_names)}"
            print(f"[warm_up_mcp] Running: {install_command}")
            
            # Execute the installation
            install_result = bash(install_command, suppress_exception=True)
            
            if install_result.return_code == 0:
                result["status"] = "success"
                result["message"] = f"Successfully installed {len(package_names)} package(s)"
                print(f"[warm_up_mcp] {result['message']}")
                print(f"[warm_up_mcp] Installed packages: {', '.join(package_names)}")
            else:
                result["status"] = "partial"
                result["message"] = f"Installation completed with warnings/errors"
                result["errors"].append(install_result.stderr)
                print(f"[warm_up_mcp] {result['message']}")
                print(f"[warm_up_mcp] stderr: {install_result.stderr}")
            
        except json.JSONDecodeError as e:
            result["status"] = "error"
            result["message"] = f"Failed to parse MCP config: {e}"
            result["errors"].append(str(e))
            print(f"[warm_up_mcp] Error: {result['message']}")
        except Exception as e:
            result["status"] = "error"
            result["message"] = f"Unexpected error: {e}"
            result["errors"].append(str(e))
            print(f"[warm_up_mcp] Error: {result['message']}")
        
        return result
