import tempfile
from typing import Literal, Optional, Dict, Any
import os
import json
import contextlib
from pathlib import Path
from blocks import bash
import time
import threading
from blocks_control_sdk.logger import log, stream_line
from blocks_control_sdk.parsers.base_messages import print_litellm_model_response, BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE
from blocks_control_sdk.parsers.cursor_messages import (
    to_model_response,
    extract_tool_info,
    is_tool_started,
)
from pydantic import BaseModel
from blocks_control_sdk.constants.cursor import CursorModels
from .agent_base import (
    CodingAgentBaseCLI,
    LLM_API_KEYS,
    LLM,
    NotifyMessageArgs,
    NotifyCompleteArgs,
    NotifyToolCallArgs,
    NotifyStartArgs,
    NotifyResumeArgs
)
from blocks_control_sdk.utils import to_pascal_case


class CursorAgentConfig(BaseModel):
    """Configuration for Cursor CLI agent."""
    model: CursorModels = CursorModels.claude_4_5_sonnet
    mode: str = "agent"  # "agent" | "plan" | "ask"


class CursorAgentCLI(CodingAgentBaseCLI):
    """
    Cursor CLI agent implementation following the Claude pattern.
    
    Uses stdout streaming via temp file + background thread for log tailing.
    Thread coordination via threading.Event() objects for clean shutdown.
    """

    def __init__(
        self,
        is_background: bool = True,
        chat_thread_id: str = None,
        config: CursorAgentConfig = None
    ):
        super().__init__(chat_thread_id)
        self.config = config or CursorAgentConfig()
        self.is_background = is_background
        self.temp_file_log = tempfile.mktemp()

        # Thread coordination events (MATCH CLAUDE PATTERN EXACTLY)
        self.log_thread_should_stop = threading.Event()
        self.log_thread_finished = threading.Event()
        self.log_thread = None
        self.log_file_position = 0  # CRITICAL: persists across thread restarts

    def _process_log_line_for_messages(self, line: str):
        """
        Parse NDJSON line and trigger notifications.

        Handles:
        - assistant messages (text content)
        - tool_call events (various tool invocations)
        - result events (skipped, handled in on_complete)
        """
        stream_line(line)

        try:
            obj = json.loads(line)

            # Handle tool_call events directly
            if obj.get("type") == "tool_call":
                self._handle_tool_call_event(obj)
                return

            # Skip result events (handled in on_complete)
            if obj.get("type") == "result":
                return

            # Skip non-assistant messages
            if obj.get("type") != "assistant":
                return

            log.debug("*" * 100)
            print_litellm_model_response(to_model_response(obj))

            # Convert to ModelResponse
            model_response = to_model_response(obj)
            message = model_response.choices[0].message
            tools = message.tool_calls or []

            # Process tool calls from message
            for tool in tools:
                log.debug("!" * 100)
                log.debug(f"<<AGENT TOOL CALL>> | {tool.function.name}")
                log.debug("!" * 100)

                try:
                    parsed_args = json.loads(tool.function.arguments)

                    # Truncate large arguments
                    if self.tool_call_arg_kb_size(tool.function.arguments) >= 10:
                        for key, value in parsed_args.items():
                            if isinstance(value, str) and self.tool_call_arg_kb_size(value) >= 10:
                                parsed_args[key] = self.tool_call_arg_kb_size_truncate_to_limit(value, 10)

                    # Handle blocks-internal-mcp tools
                    if tool.function.name and "blocks-internal-mcp" in tool.function.name:
                        _updated_name = to_pascal_case(
                            tool.function.name.replace("mcp__blocks-internal-mcp__", "")
                        )
                        parsed_args["__name__"] = _updated_name
                    else:
                        parsed_args["__name__"] = tool.function.name

                    notification = NotifyToolCallArgs(
                        tool_name=tool.function.name,
                        serialized_args=json.dumps(parsed_args),
                    )
                    self.notify_v2(notification)

                except Exception as e:
                    log.error(f"Error parsing tool {tool.function.name}: {e}")

            # Notify message if it's not a final message
            is_final = message.provider_specific_fields.get(
                BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE, False
            )

            if (
                message.content
                and message.role == "assistant"
                and not is_final
            ):
                self.assistant_messages.append(message)

                try:
                    notification = NotifyMessageArgs(message=message)
                    self.notify_v2(notification)
                except Exception as e:
                    log.error(f"Error notifying message v2: {e}")

            log.debug("*" * 100)

        except json.JSONDecodeError:
            # Skip non-JSON lines (e.g., debug output)
            pass
        except Exception as e:
            log.debug("*" * 100)
            log.error(f"Error processing log line: {e}")
            log.debug("*" * 100)

    def _handle_tool_call_event(self, obj: Dict[str, Any]):
        """
        Handle Cursor-specific tool_call events.

        Cursor emits tool_call events with subtype "started" or "completed".
        Only notify on start to avoid duplicates.

        Parsing logic is delegated to cursor_messages.extract_tool_info().
        """
        # Business logic: only notify on "started" to avoid duplicates
        if not is_tool_started(obj):
            return

        # Parser handles format-specific extraction
        tool_info = extract_tool_info(obj)
        if not tool_info:
            return

        tool_name = tool_info["tool_name"]
        args = dict(tool_info["args"])  # Copy to avoid mutating original

        # Agent handles truncation
        for key, value in list(args.items()):
            if isinstance(value, str) and self.tool_call_arg_kb_size(value) >= 10:
                args[key] = self.tool_call_arg_kb_size_truncate_to_limit(value, 10)

        args["__name__"] = tool_name

        log.debug("!" * 100)
        log.debug(f"<<AGENT TOOL CALL>> | {tool_name}")
        log.debug("!" * 100)

        try:
            self.notify_v2(NotifyToolCallArgs(
                tool_name=tool_name,
                serialized_args=json.dumps(args)
            ))
        except Exception as e:
            log.error(f"Error notifying tool call v2: {e}")

    def _start_log_thread(self):
        """Start background thread to tail temp file."""
        def tail_log_file():
            file_position = self.log_file_position  # Load from instance at start
            log.debug(f"Log thread: Starting...")

            while not self.log_thread_should_stop.is_set():
                if not os.path.exists(self.temp_file_log):
                    time.sleep(0.5)
                    continue

                try:
                    with open(self.temp_file_log, "r") as f:
                        f.seek(file_position)
                        for line in f:
                            if line.strip():
                                self._process_log_line_for_messages(line)

                        file_position = f.tell()
                        self.log_file_position = file_position  # Persist position

                except Exception as e:
                    log.error(f"Log thread: Error tailing log file: {e}")

                time.sleep(0.25)  # 250ms polling interval

            # FINAL PROCESSING ON SHUTDOWN (CRITICAL)
            time.sleep(1)  # Grace period for final writes

            if os.path.exists(self.temp_file_log):
                log.debug("Log thread: Doing final read pass...")
                try:
                    with open(self.temp_file_log, "r") as f:
                        f.seek(file_position)
                        for line in f:
                            if line.strip():
                                self._process_log_line_for_messages(line)

                        # Update final file position
                        file_position = f.tell()
                        self.log_file_position = file_position
                except Exception as e:
                    log.error(f"Log thread: Error in final read: {e}")

            # Save final position before exiting
            self.log_file_position = file_position

            # Signal that log thread has finished
            self.log_thread_finished.set()
            log.debug(f"Log thread: Finished and signaled completion (position: {file_position})")

        # Guard: Don't start another thread if one is already running
        if self.log_thread and self.log_thread.is_alive():
            return

        # Clear events for fresh start
        self.log_thread_should_stop.clear()
        self.log_thread_finished.clear()

        # Start the thread
        self.log_thread = threading.Thread(target=tail_log_file, daemon=True)
        self.log_thread.start()

    def _restart_log_thread(self, reset_position: bool = False):
        """
        Stop and restart log thread for new query.

        Args:
            reset_position: If True, reset file position to 0 (for new sessions)
        """
        if self.log_thread and self.log_thread.is_alive():
            # Stop existing thread first
            self.log_thread_should_stop.set()
            self.log_thread_finished.wait(timeout=2.0)

        # Reset status to in-progress for new query
        self.status = self.AgentStatus.TURNS_IN_PROGRESS

        # Optionally reset file position for new sessions
        if reset_position:
            self.log_file_position = 0
            log.debug("Log thread: Reset file position to 0")

            # Clean up old log file for new session
            if os.path.exists(self.temp_file_log):
                os.remove(self.temp_file_log)

        # Start new thread
        self._start_log_thread()

    def _start(self):
        """Mark session as active."""
        self.is_session_active = True

    def interrupt(self):
        """Interrupt the agent process."""
        super().interrupt()

    def _hot_reload_cleanup(self):
        """Reset log thread state for clean restart during hot reload."""
        # Stop log thread gracefully so it doesn't process stale data
        if self.log_thread and self.log_thread.is_alive():
            self.log_thread_should_stop.set()
            self.log_thread_finished.wait(timeout=2)

        # Reset interrupt flag for new query
        self.is_interrupt_triggered = False

    def write_config(self, additional_mcp_servers: Dict = None, options: Optional[dict] = None) -> str:
        """
        Write Cursor CLI configuration.

        Cursor uses CURSOR_API_KEY environment variable for authentication.
        MCP servers are written to ~/.cursor/mcp.json (per Cursor docs).
        General config (model, mode, permissions) written to ~/.cursor/config.json.
        """
        additional_mcp_servers = additional_mcp_servers or {}
        options = options or {}

        home_dir = str(Path.home())
        cursor_api_key = LLM_API_KEYS.get(LLM.CURSOR)

        if not cursor_api_key:
            log.info("Warning: CURSOR_API_KEY not found in environment")
            log.info("You may need to run 'agent login' to authenticate")

        # Read shared MCP config
        mcp_dir = os.path.join(home_dir, ".config/blocks/mcp.json")
        mcp_servers = {}

        with contextlib.suppress(Exception):
            if os.path.exists(mcp_dir):
                with open(mcp_dir, "r") as f:
                    mcp_servers = json.load(f)

        # Write MCP config to ~/.cursor/mcp.json
        cursor_config_dir = os.path.join(home_dir, ".cursor")
        os.makedirs(cursor_config_dir, exist_ok=True)

        if mcp_servers or additional_mcp_servers or self._mcp_servers:
            mcp_config_path = os.path.join(cursor_config_dir, "mcp.json")
            mcp_config_data = {
                "mcpServers": {
                    **mcp_servers,
                    **self._mcp_servers,
                    **additional_mcp_servers
                }
            }
            with open(mcp_config_path, "w") as f:
                json.dump(mcp_config_data, f, indent=2)

            log.info("*" * 100)
            log.info("CURSOR MCP CONFIG")
            log.info(json.dumps(mcp_config_data, indent=4))
            log.info("*" * 100)

        # Write general config to ~/.cursor/config.json
        config_data = {
            "model": self.config.model,
            "mode": self.config.mode,
            "permissions": {
                "allow": ["Shell(*)", "Read(*)", "Write(*)"],
                "deny": []
            },
        }

        cursor_config_path = os.path.join(cursor_config_dir, "config.json")
        with open(cursor_config_path, "w") as f:
            json.dump(config_data, f, indent=2)

        return cursor_config_path

    def query(self, query: str, tail: bool = True):
        """
        Execute Cursor agent CLI and stream output.
        
        Args:
            query: The prompt to send to the agent
            tail: Whether to tail the log file for real-time output
            
        Returns:
            BackgroundCommandOutput from the bash command
        """
        # Check resume mode BEFORE calling _start() which sets is_session_active = True
        should_resume = self.is_session_active

        # Handle interrupt flag
        is_interrupt_triggered = self.is_interrupt_triggered
        if is_interrupt_triggered:
            self.is_interrupt_triggered = False

        # Plan mode prefix
        if self.is_plan_mode_active:
            query = (
                "<system-reminder>Plan mode is active. The user indicated that they do not want you "
                "to implement the plan yet -- you MUST NOT make any edits (with the exception of the "
                "plan file mentioned below), run any non-readonly tools (including changing configs "
                "or making commits), or otherwise make any changes to the system.</system-reminder> "
                + query
            )

        # Restart log thread if needed for new query (only if tailing is enabled)
        if tail:
            # Reset position only for new sessions, not for resume
            self._restart_log_thread(reset_position=not should_resume)

        # Write query to temp file
        temp_file_path = tempfile.mktemp()
        with open(temp_file_path, "w") as f:
            f.write(query)
            f.flush()

        # Build command args
        args = [
            "--approve-mcps",
            "--force",
            "--output-format", "stream-json",
            "--print",
        ]

        # Add mode if not default agent mode
        if self.config.mode != "agent":
            args.extend(["--mode", self.config.mode])

        # Override mode for plan mode
        if self.is_plan_mode_active:
            args.extend(["--mode", "plan"])

        # Add resume flag if continuing session
        if should_resume:
            log.debug("-" * 100)
            log.info("Resuming session")
            log.debug("-" * 100)
            args.append("--continue")

            try:
                notification = NotifyResumeArgs()
                self.notify_v2(notification)
            except Exception as e:
                log.error(f"Error notifying resume v2: {e}")

        # Debug: print the query being sent
        log.debug("-" * 100)
        log.debug(f"Contents of file {temp_file_path}:")
        with open(temp_file_path, "r") as f:
            log.debug(f.read())
        log.debug("-" * 100)

        temp_file_log = self.temp_file_log

        # Set status and notify start
        self.status = self.AgentStatus.TURNS_IN_PROGRESS

        try:
            notification = NotifyStartArgs()
            self.notify_v2(notification)
        except Exception as e:
            log.error(f"Error notifying start v2: {e}")

        # Execute command
        # Note: Using 'agent' as the CLI command name for Cursor
        agent_command = f"cat {temp_file_path} | agent {' '.join(args)} >> {temp_file_log} 2>&1"

        log.info("*" * 100)
        log.info(f"Running cursor command: {agent_command}")
        log.info("*" * 100)

        out = bash(
            agent_command,
            background=self.is_background,
            blocklist_prefixes=["BLOCKS_", "AWS_", "ECS_", "E2B"]
        )

        def on_complete(cmd_output):
            """Callback when the agent command completes."""
            stdout_output = cmd_output.stdout
            stderr_output = cmd_output.stderr
            pid = cmd_output.pid
            return_code = cmd_output.return_code

            # First interrupt check
            if self.is_interrupt_triggered:
                return

            self.status = self.AgentStatus.TURNS_COMPLETED

            log.info(f"Cursor agent subprocess {pid} has exited with code {return_code}")
            log.info(f"Cursor agent subprocess err output: {stderr_output}")
            log.info(f"Cursor agent subprocess output: {stdout_output}")

            # Give time for final messages to be written to file
            time.sleep(0.5)

            # Signal log thread to stop and wait for it to finish processing
            log.debug("Signaling log thread to stop...")
            self.log_thread_should_stop.set()

            # Wait for log thread to finish with a timeout
            if self.log_thread_finished.wait(timeout=15.0):
                log.debug("Log thread finished successfully")
            else:
                log.error("Warning: Log thread did not finish within timeout")

            # Second interrupt check before notify
            if self.is_interrupt_triggered:
                return

            # Get last message
            last_assistant_message = self.assistant_messages[-1] if self.assistant_messages else None
            last_assistant_message_content = last_assistant_message.content if last_assistant_message else ""

            log.debug(">>>Before notify complete!!!")

            try:
                notification = NotifyCompleteArgs(
                    last_message=last_assistant_message_content,
                )
                self.notify_v2(notification)
            except Exception as e:
                log.error(f"Error notifying complete v2: {e}")

            log.debug(">>>After notify complete!!!")

            # Cleanup temp file
            if os.path.exists(temp_file_path):
                os.remove(temp_file_path)

        out.register_callback(on_complete)

        self.pid = out.pid
        self.is_session_active = True
        self._set_background_command_output(out)

        return out
