"""
Kimi Code Agent CLI implementation.

This module provides the KimiAgentCLI class that wraps the Kimi CLI
for use in the Blocks agent system.
"""

import os
import tempfile
import json
import contextlib
from typing import Optional, Dict, Any
from pathlib import Path
from blocks import bash
import time
import threading
from blocks_control_sdk.parsers.base_messages import print_litellm_model_response
from blocks_control_sdk.parsers.kimi_messages import (
    to_model_response,
    extract_tool_info,
    is_assistant_message,
    is_tool_result,
    get_text_content,
    KIMI_TOOL_NAMES,
)
from pydantic import BaseModel
from blocks_control_sdk.constants.kimi import KimiModels
from .agent_base import (
    CodingAgentBaseCLI,
    LLM_API_KEYS,
    LLM,
    NotifyMessageArgs,
    NotifyCompleteArgs,
    NotifyToolCallArgs,
    NotifyStartArgs,
    NotifyResumeArgs,
)
from blocks_control_sdk.utils import to_pascal_case, is_bash_error_code, render_agent_error_message
from blocks_control_sdk.logger import log, stream_line


class KimiAgentConfig(BaseModel):
    """Configuration for Kimi CLI agent."""
    model: KimiModels = KimiModels.kimi_default


class KimiAgentCLI(CodingAgentBaseCLI):
    """
    Kimi CLI agent implementation following the Cursor/OpenCode pattern.

    Uses --print --output-format=stream-json for non-interactive streaming output.
    Output is captured to a temp file and tailed by a background thread.
    """

    def __init__(
        self,
        is_background: bool = True,
        chat_thread_id: str = None,
        config: KimiAgentConfig = None,
    ):
        super().__init__(chat_thread_id)
        self.config = config or KimiAgentConfig()
        self.is_background = is_background
        self.temp_file_log = tempfile.mktemp(suffix=".jsonl", prefix="kimi_")

        # Thread coordination events
        self.log_thread_should_stop = threading.Event()
        self.log_thread_finished = threading.Event()
        self.log_thread = None
        self.log_file_position = 0  # Persists across thread restarts

        # Write read-only explore agent config
        self.explore_agent_file = self._write_explore_agent_config()

    def _write_explore_agent_config(self) -> str:
        """Write a read-only Kimi agent config that excludes mutating tools."""
        config_dir = os.path.join(str(Path.home()), ".config", "blocks")
        os.makedirs(config_dir, exist_ok=True)
        config_path = os.path.join(config_dir, "kimi_explore.config.yaml")

        if os.path.exists(config_path):
            return config_path

        yaml_content = """version: 1
agent:
  extend: default
  exclude_tools:
    - "kimi_cli.tools.shell:Shell"
    - "kimi_cli.tools.file:WriteFile"
    - "kimi_cli.tools.file:StrReplaceFile"
"""
        with open(config_path, "w") as f:
            f.write(yaml_content)

        return config_path

    def _process_log_line(self, line: str):
        """
        Parse a JSONL line from Kimi's stream-json output and trigger notifications.

        Handles:
        - assistant messages with text content
        - assistant messages with tool_calls
        - tool result messages (skipped, handled internally by Kimi)
        """
        stream_line(line)

        try:
            obj = json.loads(line)

            # Skip tool result messages (role=tool)
            if is_tool_result(obj):
                return

            # Only process assistant messages
            if not is_assistant_message(obj):
                return

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

            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"[Kimi] Error parsing tool {tool.function.name}: {e}")

            # Store and notify assistant messages with content
            if message.content and message.role == "assistant":
                self.assistant_messages.append(message)
                try:
                    notification = NotifyMessageArgs(message=message)
                    self.notify_v2(notification)
                except Exception as e:
                    log.error(f"[Kimi] Error notifying message: {e}")

            log.debug("*" * 100)

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

    def _start_log_thread(self):
        """Start background thread to tail temp file."""
        def tail_log_file():
            file_position = self.log_file_position
            log.debug("[Kimi] 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(line)

                        file_position = f.tell()
                        self.log_file_position = file_position

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

                time.sleep(0.25)

            # Final processing on shutdown
            time.sleep(1)

            if os.path.exists(self.temp_file_log):
                log.debug("[Kimi] 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(line)

                        file_position = f.tell()
                        self.log_file_position = file_position
                except Exception as e:
                    log.error(f"[Kimi] Log thread: Error in final read: {e}")

            self.log_file_position = file_position
            self.log_thread_finished.set()
            log.debug(f"[Kimi] Log thread: Finished (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, name="kimi-log-tail")
        self.log_thread.start()

    def _restart_log_thread(self, reset_position: bool = False):
        """
        Stop and restart log thread for a 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():
            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

        if reset_position:
            self.log_file_position = 0
            log.debug("[Kimi] 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)

        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."""
        if self.log_thread and self.log_thread.is_alive():
            self.log_thread_should_stop.set()
            self.log_thread_finished.wait(timeout=2)

        self.is_interrupt_triggered = False

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

        Kimi uses ~/.kimi/config.toml for provider/model configuration
        and ~/.kimi/mcp.json for MCP server configuration.
        """
        additional_mcp_servers = additional_mcp_servers or {}
        options = options or {}

        home_dir = str(Path.home())
        kimi_api_key = LLM_API_KEYS.get(LLM.KIMI)

        if not kimi_api_key:
            log.info("[Kimi] Warning: KIMI_API_KEY not found in environment")

        # 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)

        # Ensure ~/.kimi directory exists
        kimi_config_dir = os.path.join(home_dir, ".kimi")
        os.makedirs(kimi_config_dir, exist_ok=True)

        # Write MCP config to ~/.kimi/mcp.json
        if mcp_servers or additional_mcp_servers or self._mcp_servers:
            mcp_config_path = os.path.join(kimi_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("KIMI MCP CONFIG")
            log.info(json.dumps(mcp_config_data, indent=4))
            log.info("*" * 100)

        # Write config.toml with provider and model configuration
        config_toml_path = os.path.join(kimi_config_dir, "config.toml")

        model_name = self.config.model.value

        # Build TOML content using moonshot-ai provider (correct configuration)
        toml_lines = [
            '# Kimi CLI configuration (auto-generated by Blocks)',
            '',
            f'default_model = "{model_name}"',
            'default_thinking = true',
            'default_yolo = true',
            '',
            # Model definitions for moonshot-ai provider
            '[models."moonshot-ai/kimi-k2-0711-preview"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2-0711-preview"',
            'max_context_size = 131072',
            '',
            '[models."moonshot-ai/kimi-k2-0905-preview"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2-0905-preview"',
            'max_context_size = 262144',
            '',
            '[models."moonshot-ai/kimi-k2.5"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2.5"',
            'max_context_size = 262144',
            'capabilities = ["thinking", "image_in", "video_in"]',
            '',
            '[models."moonshot-ai/kimi-k2-thinking"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2-thinking"',
            'max_context_size = 262144',
            'capabilities = ["thinking", "always_thinking"]',
            '',
            '[models."moonshot-ai/kimi-k2-turbo-preview"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2-turbo-preview"',
            'max_context_size = 262144',
            '',
            '[models."moonshot-ai/kimi-k2-thinking-turbo"]',
            'provider = "managed:moonshot-ai"',
            'model = "kimi-k2-thinking-turbo"',
            'max_context_size = 262144',
            'capabilities = ["thinking", "always_thinking"]',
            '',
            # Provider configuration
            '[providers."managed:moonshot-ai"]',
            'type = "kimi"',
            'base_url = "https://api.moonshot.ai/v1"',
        ]

        if kimi_api_key:
            toml_lines.append(f'api_key = "{kimi_api_key}"')

        toml_lines.extend([
            '',
            '[loop_control]',
            'max_steps_per_turn = 100',
            'max_retries_per_step = 3',
            'max_ralph_iterations = 0',
            'reserved_context_size = 50000',
            '',
            '[services]',
            '',
            '[mcp.client]',
            'tool_call_timeout_ms = 60000',
        ])

        toml_content = "\n".join(toml_lines) + "\n"

        with open(config_toml_path, "w") as f:
            f.write(toml_content)

        log.info("*" * 100)
        log.info("KIMI CONFIG TOML")
        log.info(toml_content)
        log.info("*" * 100)

        return config_toml_path

    def query(self, query: str, tail: bool = True):
        """
        Execute Kimi 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
        if tail:
            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 = [
            "--print",
            "--output-format=stream-json",
        ]

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

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

        # Debug: print the query being sent
        log.debug("-" * 100)
        log.debug(f"[Kimi] 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"[Kimi] Error notifying start: {e}")

        # Execute command - pass prompt via -p flag (stdin piping hangs with kimi CLI)
        agent_command = f"kimi {' '.join(args)} -p \"$(cat {temp_file_path})\" >> {temp_file_log} 2>&1"

        log.debug("*" * 100)
        log.debug(f"[Kimi] Running command: {agent_command}")
        log.debug("*" * 100)

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

        def on_complete(cmd_output):
            """Callback when the Kimi agent command completes."""
            print("[Kimi] on_complete")
            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"[Kimi] Process {pid} exited with code {return_code}")
            log.info(f"[Kimi] stderr: {stderr_output}")
            log.info(f"[Kimi] stdout: {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
            log.debug("[Kimi] Signaling log thread to stop...")
            self.log_thread_should_stop.set()

            if self.log_thread_finished.wait(timeout=15.0):
                log.debug("[Kimi] Log thread finished successfully")
            else:
                log.debug("[Kimi] 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 ""

            if is_bash_error_code(return_code):
                last_assistant_message_content = render_agent_error_message(stderr_output, last_assistant_message_content)

            log.debug("[Kimi] Before notify complete")

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

            log.debug("[Kimi] 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
