"""
OpenCode Agent CLI implementation.

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

import os
import tempfile
import json
from typing import List, Optional
from blocks import bash
from pathlib import Path
import threading
import time
from blocks_control_sdk.logger import log
from blocks_control_sdk.utils import is_bash_error_code, render_agent_error_message, to_pascal_case
from blocks_control_sdk.parsers.opencode_messages import (
    to_model_response,
    extract_session_id,
    extract_tool_info,
    is_text_event,
    is_tool_event,
    is_error_event,
    get_text_content,
    get_error_message,
    OPENCODE_TOOL_NAMES,
)
from blocks_control_sdk.parsers.base_messages import print_litellm_model_response
from pydantic import BaseModel
from .agent_base import (
    CodingAgentBaseCLI,
    NotifyCompleteArgs,
    NotifyMessageArgs,
    NotifyResumeArgs,
    NotifyStartArgs,
    NotifyToolCallArgs,
)
from blocks_control_sdk.constants.opencode import OpenCodeModels, SisyphusModels


class OpenCodeAgentConfig(BaseModel):
    """Configuration for OpenCode agent."""
    model: OpenCodeModels = OpenCodeModels.claude_sonnet_latest

class SisyphusAgentConfig(BaseModel):
    """Configuration for Sisyphus agent."""
    model: SisyphusModels = SisyphusModels.sisyphus_latest


class OpenCodeAgentCLI(CodingAgentBaseCLI):
    """
    OpenCode CLI agent implementation.

    Wraps the OpenCode CLI (opencode run) to provide LLM-backed coding assistance
    through the Blocks agent system.
    """

    def __init__(self, chat_thread_id: str = None, config: OpenCodeAgentConfig = OpenCodeAgentConfig()):
        super().__init__(chat_thread_id)
        self.config = config
        self._temp_log_file: Optional[str] = None
        self._log_file_position = 0
        self._tracked_session_ids: List[str] = []
        self._session_ids_lock = threading.Lock()
        self._log_thread_stop = threading.Event()
        self._log_thread_finished = threading.Event()

        # Create temp log file for capturing JSON output
        self._temp_log_file = tempfile.mktemp(suffix=".jsonl", prefix="opencode_")

        self.log_thread = threading.Thread(
            target=self._tail_log_file, daemon=True, name="opencode-log-tail"
        )
        self.log_thread.start()

    def _tail_log_file(self):
        """Background thread to tail the JSON log file and process events."""
        while not self._log_thread_stop.is_set():
            if not self._temp_log_file or not os.path.exists(self._temp_log_file):
                log.debug(f"Log file not found: {self._temp_log_file}")
                time.sleep(0.5)
                continue

            try:
                with open(self._temp_log_file, "r") as f:
                    f.seek(self._log_file_position)
                    for line in f:
                        line = line.strip()
                        if not line:
                            log.debug(f"Empty line: {line}")
                            continue

                        log.debug(line)

                        try:
                            event = json.loads(line)
                            self._process_event(event)
                        except json.JSONDecodeError as e:
                            log.error(f"Error parsing JSON line: {e}")
                        except Exception as e:
                            log.error(f"Error processing event: {e}")

                    self._log_file_position = f.tell()
            except FileNotFoundError:
                self._log_file_position = 0
            except Exception as e:
                log.error(f"Error reading log file: {e}")

            time.sleep(0.5)

        # Final read pass to catch any remaining messages after stop signal
        if self._temp_log_file and os.path.exists(self._temp_log_file):
            try:
                with open(self._temp_log_file, "r") as f:
                    f.seek(self._log_file_position)
                    for line in f:
                        line = line.strip()
                        if not line:
                            continue

                        log.debug(line)

                        try:
                            event = json.loads(line)
                            self._process_event(event)
                        except json.JSONDecodeError as e:
                            log.error(f"Error parsing JSON line (final pass): {e}")
                        except Exception as e:
                            log.error(f"Error processing event (final pass): {e}")

                    self._log_file_position = f.tell()
            except Exception as e:
                log.error(f"Error in final read pass: {e}")

        self._log_thread_finished.set()

    def _process_event(self, event: dict):
        """
        Process a single OpenCode JSON event.

        Args:
            event: Parsed JSON event dict
        """
        event_type = event.get("type")

        # Extract and track session ID
        session_id = extract_session_id(event)
        if session_id:
            with self._session_ids_lock:
                if session_id not in self._tracked_session_ids:
                    self._tracked_session_ids.append(session_id)
                    log.info(f"Tracking session ID: {session_id}")

        log.debug("*" * 100)
        log.debug(f"Processing event type: {event_type}")

        try:
            model_response = to_model_response(event)
            print_litellm_model_response(model_response)

            message = model_response.choices[0].message
            tools = message.tool_calls or []

            if is_tool_event(event):
                # Handle tool use event
                tool_info = extract_tool_info(event)
                if tool_info:
                    tool_name = tool_info["tool_name"]
                    tool_input = tool_info["input"]

                    # Check if tool is completed
                    if tool_info.get("status") == "completed":
                        try:
                            # Truncate large arguments
                            serialized_input = json.dumps(tool_input)
                            if self.tool_call_arg_kb_size(serialized_input) >= 10:
                                for key, value in tool_input.items():
                                    if isinstance(value, str) and self.tool_call_arg_kb_size(value) >= 10:
                                        tool_input[key] = self.tool_call_arg_kb_size_truncate_to_limit(value, 10)

                            # Format tool name for notification
                            display_name = OPENCODE_TOOL_NAMES.get(tool_name, tool_name)

                            # Handle internal MCP tools
                            if "blocks-internal-mcp" in tool_name:
                                display_name = to_pascal_case(tool_name.replace("mcp__blocks-internal-mcp__", ""))
                                tool_input["__name__"] = display_name

                            notification = NotifyToolCallArgs(
                                tool_name=tool_name,
                                serialized_args=json.dumps({
                                    "__name__": display_name,
                                    **tool_input,
                                }),
                            )
                            self.notify_v2(notification)
                        except Exception as e:
                            log.error(f"Error notifying tool call: {e}")

            elif is_text_event(event):
                # Handle text message event
                text_content = get_text_content(event)
                if text_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"Error notifying message: {e}")

            elif is_error_event(event):
                # Handle error event
                error_message = get_error_message(event)
                log.error(f"Error event: {error_message}")

        except ValueError as e:
            # Unknown event type
            log.debug(f"Skipping event: {e}")

        log.debug("*" * 100)

    def _start(self):
        """Mark the agent as started."""
        pass

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

    def _restart_log_thread(self, reset_position=False):
        """Restart the log thread for a new query.

        Args:
            reset_position: If True, reset file position to 0 (for new sessions).
                           If False, preserve position to continue from last read (for resumes).
        """
        # Stop existing thread if alive
        if self.log_thread and self.log_thread.is_alive():
            self._log_thread_stop.set()
            self._log_thread_finished.wait(timeout=2)

        # Clear events for fresh start
        self._log_thread_stop.clear()
        self._log_thread_finished.clear()

        # Only reset position for new sessions
        if reset_position:
            self._log_file_position = 0

        # Start new thread
        self.log_thread = threading.Thread(
            target=self._tail_log_file, daemon=True, name="opencode-log-tail"
        )
        self.log_thread.start()

    def _hot_reload_cleanup(self):
        """
        Reset state for hot reload.

        NOTE: Do NOT clear _tracked_session_ids - we need them for session continuity.
        NOTE: _log_file_position is preserved across resumes to avoid replaying old messages.
        """
        # Stop log thread gracefully so it doesn't process stale data
        if self.log_thread and self.log_thread.is_alive():
            self._log_thread_stop.set()
            self._log_thread_finished.wait(timeout=2)

        self.is_interrupt_triggered = False

    def query(self, query: str, tail: bool = True):
        """
        Execute a query using the OpenCode CLI.

        Args:
            query: The prompt/query to send to OpenCode
            tail: Whether to tail the output (default True)

        Returns:
            BackgroundCommandOutput from the bash command
        """

        if os.getenv("OMOC_INSTALLED") == "1":
            log.info("-" * 100)
            log.info("Sisyphys is installed")
            log.info("-" * 100)
            # blocks until /tmp/omoc.lock exists
            while not os.path.exists("/tmp/omoc.lock"):
                time.sleep(0.1)
            log.info("Sisyphys is ready")
            log.info("-" * 100)

        self._start()
        is_interrupt_triggered = self.is_interrupt_triggered
        if is_interrupt_triggered:
            self.is_interrupt_triggered = False

        # Add plan mode reminder if active
        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), run any non-readonly tools (including changing "
                "configs or making commits), or otherwise make any changes to the system."
                "</system-reminder> " + query
            )

        should_resume = self.is_session_active

        # Clear tracked session IDs on new session (not resume)
        if not should_resume:
            with self._session_ids_lock:
                self._tracked_session_ids.clear()

        # Restart log thread for every query (only if tailing is enabled)
        # Match Claude: never reset position - just continue from where we left off
        if tail:
            self._restart_log_thread()

        # Create temp file for the query (to handle long prompts)
        temp_query_file = tempfile.mktemp()
        with open(temp_query_file, 'w') as f:
            f.write(query)

        # Build command arguments
        args = [
            "--format", "json",
        ]

        # Add model if specified
        if self.config.model:
            args.extend(["--model", self.config.model.value])

        # Handle session resume
        if should_resume and self._tracked_session_ids:
            last_session_id = self._tracked_session_ids[-1]
            log.info("-" * 100)
            log.info(f"Resuming session: {last_session_id}")
            log.info("-" * 100)

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

            args.extend(["--session", last_session_id])

        # Add the query via cat
        args.append(f'"$(cat {temp_query_file})"')

        # Update 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: {e}")

        # Build and execute the command
        # Append output to log file for the tail thread (>> preserves position tracking)
        cmd = f"opencode run {' '.join(args)} >> {self._temp_log_file} 2>&1"
        log.info(f"Executing: {cmd}")

        out = bash(
            cmd,
            background=True,
            blocklist_prefixes=["BLOCKS_", "AWS_", "ECS_", "E2B"],
        )

        def on_complete(cmd_output):
            """Callback when the OpenCode process completes."""
            if self.is_interrupt_triggered:
                return

            self.status = self.AgentStatus.TURNS_COMPLETED
            stdout_output = cmd_output.stdout
            stderr_output = cmd_output.stderr
            pid = cmd_output.pid
            return_code = cmd_output.return_code

            log.info(f"Process {pid} exited with code {return_code}")
            if stderr_output:
                log.error(f"stderr: {stderr_output}")

            # Brief grace period for final writes, then signal log thread to stop
            time.sleep(0.5)
            self._log_thread_stop.set()
            if not self._log_thread_finished.wait(timeout=15.0):
                log.error("Warning: Log thread did not finish within timeout")

            # Get last assistant 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.info(f"Last message (is_bash_error: {is_bash_error_code(return_code)}): {last_assistant_message_content[:200] if last_assistant_message_content else 'None'}...")

            # Handle error cases
            if is_bash_error_code(return_code):
                # Try to read error from log file
                error_from_log = None
                if self._temp_log_file and os.path.exists(self._temp_log_file):
                    try:
                        with open(self._temp_log_file, 'r') as f:
                            lines = f.readlines()
                            # Check last few lines for error events
                            for line in lines[-10:]:
                                try:
                                    event = json.loads(line.strip())
                                    if event.get("type") == "error":
                                        error_from_log = get_error_message(event)
                                        break
                                except:
                                    pass
                    except:
                        pass

                error_content = error_from_log or last_assistant_message_content

                # Handle common auth errors
                if error_content and ("401" in error_content or "Unauthorized" in error_content):
                    error_content = "Error: Authentication failed. Please check your API key is valid."

                last_assistant_message_content = render_agent_error_message(stderr_output, error_content)

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

        out.register_callback(on_complete)

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

        return out

    def write_config(self, additional_mcp_servers={}, options=None) -> str:
        """
        Write OpenCode configuration files.

        Writes auth credentials to ~/.local/share/opencode/auth.json
        and MCP configuration.

        Args:
            additional_mcp_servers: Additional MCP servers to configure
            options: Additional options dict

        Returns:
            Path to the auth config file
        """
        auth_dir = Path.home() / ".local" / "share" / "opencode"
        auth_file = auth_dir / "auth.json"
        config_dir = Path.home() / ".config" / "opencode"
        config_file = config_dir / "opencode.json"  # Must be opencode.json, not config.json

        # Ensure directories exist
        auth_dir.mkdir(parents=True, exist_ok=True)
        config_dir.mkdir(parents=True, exist_ok=True)

        # Build auth config
        auth_config = {}

        # Add API keys from environment
        anthropic_key = os.getenv("ANTHROPIC_API_KEY")
        openai_key = os.getenv("OPENAI_API_KEY")
        gemini_key = os.getenv("GEMINI_API_KEY")

        if anthropic_key:
            auth_config["anthropic"] = {"apiKey": anthropic_key}
        if openai_key:
            auth_config["openai"] = {"apiKey": openai_key}
        if gemini_key:
            auth_config["google"] = {"apiKey": gemini_key}

        # Check if existing auth has OAuth tokens and preserve them
        if auth_file.exists():
            try:
                with open(auth_file, "r") as f:
                    existing_auth = json.load(f)
                # Merge: keep existing OAuth tokens, update API keys
                for provider, config in existing_auth.items():
                    if provider not in auth_config:
                        auth_config[provider] = config
                    elif "oauth" in config or "accessToken" in config:
                        # Preserve OAuth data
                        auth_config[provider].update(config)
            except Exception as e:
                log.error(f"Error reading existing auth: {e}")

        # Write auth config
        try:
            with open(auth_file, "w") as f:
                json.dump(auth_config, f, indent=2)
            log.info(f"Auth config written to {auth_file}")
        except Exception as e:
            log.error(f"Error writing auth config: {e}")

        # Build main config
        blocks_mcp_config_path = Path.home() / ".config" / "blocks" / "mcp.json"

        # Load existing config if it exists
        existing_config = {}
        if config_file.exists():
            try:
                with open(config_file, "r") as f:
                    existing_config = json.load(f)
                log.info(f"Loaded existing config from {config_file}")
            except Exception as e:
                log.error(f"Error reading existing config: {e}")

        main_config = {
            "$schema": "https://opencode.ai/config.json",
            "model": self.config.model.value if os.getenv("OMOC_INSTALLED") != "1" else OpenCodeModels.claude_opus_latest.value,
            # Allow all operations without prompting for approval
            "permission": {
                "edit": "allow",
                "write": "allow",
                "bash": "allow",
                "read": "allow",
                "glob": "allow",
                "grep": "allow",
                "fetch": "allow",
                "mcp": "allow",
                "external_directory": "allow",
                "list": "allow",
                "task": "allow",
                "skill": "allow",
                "todoread": "allow",
                "todowrite": "allow",
                "question": "deny"
            },
        }

        # Merge existing config: our values always take precedence, but preserve user's custom keys
        for key, value in existing_config.items():
            if key not in main_config:
                # Preserve custom keys that we don't explicitly set
                main_config[key] = value
            elif key == "permission" and isinstance(value, dict):
                # Only add user's custom permissions that we don't define; our permissions are never overwritten
                for perm_key, perm_value in value.items():
                    if perm_key not in main_config["permission"]:
                        main_config["permission"][perm_key] = perm_value

        # Load and convert MCP configuration from Blocks format to OpenCode format
        mcp_servers = {}
        try:
            if blocks_mcp_config_path.exists():
                with open(blocks_mcp_config_path, "r") as f:
                    blocks_mcp = json.load(f)

                # Convert Blocks MCP format to OpenCode MCP format
                for server_name, server_config in blocks_mcp.items():
                    opencode_server = {
                        "type": "local",
                        "enabled": not server_config.get("disabled", False),
                    }

                    # Convert command + args to command array
                    command = server_config.get("command", "")
                    args = server_config.get("args", [])
                    if command:
                        opencode_server["command"] = [command] + args

                    # Convert env to environment
                    if "env" in server_config:
                        opencode_server["environment"] = server_config["env"]

                    mcp_servers[server_name] = opencode_server
        except Exception as e:
            log.error(f"Error loading Blocks MCP config: {e}")

        # Add registered MCP servers (including blocks-internal-mcp)
        for server_name, server_config in self._mcp_servers.items():
            opencode_server = {
                "type": "local",
                "enabled": not server_config.get("disabled", False),
            }
            command = server_config.get("command", "")
            srv_args = server_config.get("args", [])
            if command:
                opencode_server["command"] = [command] + srv_args
            if "env" in server_config:
                opencode_server["environment"] = server_config["env"]
            mcp_servers[server_name] = opencode_server

        # Merge with additional MCP servers (also convert format)
        for server_name, server_config in additional_mcp_servers.items():
            if "command" in server_config and isinstance(server_config["command"], str):
                # Convert Blocks format to OpenCode format
                opencode_server = {
                    "type": "local",
                    "enabled": not server_config.get("disabled", False),
                    "command": [server_config["command"]] + server_config.get("args", []),
                }
                if "env" in server_config:
                    opencode_server["environment"] = server_config["env"]
                mcp_servers[server_name] = opencode_server
            else:
                # Already in OpenCode format
                mcp_servers[server_name] = server_config

        # Add MCP config if any servers configured
        if mcp_servers:
            main_config["mcp"] = mcp_servers

        # Write main config
        try:
            with open(config_file, "w") as f:
                json.dump(main_config, f, indent=2)
            log.info(f"Config written to {config_file}")
        except Exception as e:
            log.error(f"Error writing config: {e}")

        log.debug("*" * 100)
        log.debug("OPENCODE AUTH CONFIG")
        log.debug(auth_file.read_text() if auth_file.exists() else "No auth config")
        log.debug("*" * 100)
        log.debug("OPENCODE CONFIG")
        log.debug(config_file.read_text() if config_file.exists() else "No config")
        log.debug("*" * 100)

        return str(auth_file)

    def __del__(self):
        """Cleanup on destruction."""
        self._log_thread_stop.set()
        # Clean up temp log file
        if self._temp_log_file and os.path.exists(self._temp_log_file):
            try:
                os.remove(self._temp_log_file)
            except:
                pass
