import tempfile
from typing import Literal, Optional
import os
import json
import contextlib
from pathlib import Path
from blocks import bash
import json
from blocks_control_sdk.utils.main import is_bash_error_code, render_agent_error_message
from blocks_control_sdk.utils.logger import log, is_debug
import time
import threading
from blocks_control_sdk.parsers.base_messages import print_litellm_model_response, BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE
from blocks_control_sdk.parsers.anthropic_messages import (
    ANTHROPIC_TOOLS__TODOWRITE,
    ANTHROPIC_TOOLS__WRITE,
    ANTHROPIC_TOOLS__GREP,
    ANTHROPIC_TOOLS__LS,
    ANTHROPIC_TOOLS__GLOB,
    ANTHROPIC_TOOLS__BASH,
    ANTHROPIC_TOOLS__READ,
    ANTHROPIC_TOOLS__TODOWRITE,
    to_model_response,
    AnthropicTodoListUpdate,
    AnthropicTodo
)
from pydantic import BaseModel
from .agent_base import CodingAgentBaseCLI, LLM_API_KEYS, LLM, NotifyMessageArgs, NotifyCompleteArgs, NotifyToolCallArgs, NotifyStartArgs, NotifyResumeArgs
from blocks_control_sdk.constants.anthropic import AnthropicModels

class ClaudeAgentConfig(BaseModel):
    model: AnthropicModels = AnthropicModels.sonnet

class ClaudeCode(CodingAgentBaseCLI):

    auth_mode: Literal["api_key", "oauth"] = "api_key"

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

        # 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  # Persist file position across thread restarts

    def _process_log_line_for_messages(self, line):
        """Process a single log line for messages and notifications"""
        if is_debug():
            print(line, end="", flush=True)

        # Parse raw JSON first
        raw_json = None
        try:
            raw_json = json.loads(line)
        except json.JSONDecodeError as e:
            if is_debug():
                log.error(f"Error parsing JSON: {e}")
            return

        # Try to convert to ModelResponse - this is what we care about for regression testing
        try:
            log.debug("*"*100)
            print_litellm_model_response(to_model_response(raw_json))
            model_response = to_model_response(raw_json)
            message = model_response.choices[0].message
            tools = message.tool_calls or []

            # Determine parsed type for raw capture
            if tools:
                parsed_type = "tool_call"
            elif message.content and message.role == "assistant":
                parsed_type = "message"
            else:
                parsed_type = "other"

            # Capture the raw JSON with successful to_model_response parse info
            self.capture_raw_line(
                raw_json=raw_json,
                parsed_type=parsed_type,
                parse_success=True
            )

            for tool in tools:
                log.debug("!"*100)
                log.debug(f"<<AGENT TOOL CALL>> | {tool.function.name}")
                log.debug("!"*100)
                try:
                    if tool.function.name == ANTHROPIC_TOOLS__TODOWRITE:
                        parsed_args = json.loads(tool.function.arguments)
                        antropic_todo = AnthropicTodoListUpdate(**parsed_args)
                        parsed_args = {
                            "__name__": tool.function.name,
                            "__markdown__": str(antropic_todo),
                            **parsed_args,
                        }
                        serialized_args = json.dumps(parsed_args)

                        try:
                            notification = NotifyToolCallArgs(
                                tool_name=tool.function.name,
                                serialized_args=serialized_args,
                            )
                            self.notify(notification)
                        except Exception as e:
                            log.error(f"Error notifying tool call v2: {e}")

                    elif tool.function.name == ANTHROPIC_TOOLS__WRITE:
                        parsed_args = json.loads(tool.function.arguments)
                        content = parsed_args.get("content")
                        should_truncate = self.tool_call_arg_kb_size(content) >= 10
                        if should_truncate:
                            content = self.tool_call_arg_kb_size_truncate_to_limit(content, 10)
                            parsed_args["content"] = content
                        parsed_args = {
                            "__name__": tool.function.name,
                            **parsed_args,
                            "is_truncated": should_truncate,
                        }
                        serialized_args = json.dumps(parsed_args)

                        try:
                            notification = NotifyToolCallArgs(
                                tool_name=tool.function.name,
                                serialized_args=serialized_args,
                            )
                            self.notify(notification)
                        except Exception as e:
                            log.error(f"Error notifying tool call v2: {e}")
                    else:
                        try:
                            parsed_args = json.loads(tool.function.arguments)
                            # if too large
                            if self.tool_call_arg_kb_size(tool.function.arguments) >= 10:
                                # for each item in parsed_args, truncate the value if it is too large
                                for key, value in parsed_args.items():
                                    if self.tool_call_arg_kb_size(value) >= 10:
                                        parsed_args[key] = self.tool_call_arg_kb_size_truncate_to_limit(value, 10)
                            # if tool name contains "blocks-internal-mcp"

                            if "blocks-internal-mcp" in tool.function.name:
                                parsed_args = {
                                    "__name__": "blocks",
                                    "message": "Performing core blocks function."
                                }

                            notification = NotifyToolCallArgs(
                                tool_name=tool.function.name,
                                serialized_args=json.dumps({
                                    "__name__": tool.function.name,
                                    **parsed_args,
                                }),
                            )
                            self.notify(notification)
                        except Exception as e:
                            log.error(f"Error notifying tool call v2: {e}")
                except Exception as e:
                    log.error(f"Error parsing tool {tool.function.name}: {e}")

            if (
                message.content
                and message.role == "assistant"
                and not message.provider_specific_fields.get(BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE, False)
            ):
                self.assistant_messages.append(message)

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

            log.debug("*"*100)
        except Exception as e:
            # Capture parse/processing failures with the raw JSON
            if raw_json is not None:
                self.capture_raw_line(
                    raw_json=raw_json,
                    parsed_type="parse_error",
                    parse_success=False,
                    parse_error=str(e)
                )
            if is_debug():
                log.error("*"*100)
                log.error(f"Error parsing line: {e}")
                log.error("*"*100)
                

    def _start_log_thread(self):
        """Start the log tailing thread"""
        def tail_log_file():
            file_position = self.log_file_position  
            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:
                            self._process_log_line_for_messages(line)

                        file_position = f.tell()
                        self.log_file_position = file_position  # Save for next restart

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

            time.sleep(1)
            # Do one final read to catch any remaining messages
            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:
                            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})")

        if self.log_thread and self.log_thread.is_alive():
            # Log thread already running, don't start a new one
            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=False):
        """Restart the log thread if it's not running

        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)

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

        # Start new thread
        self._start_log_thread()

    def _start(self):
        pass

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

    def write_config(self, additional_mcp_servers={}, options: Optional[dict] = {}) -> str:
        """
        Write the claude config file if it doesn't exist.
        """
        options = options or {}
        use_anthropic_api_key = options.get("use_anthropic_api_key", True)

        filename = ".claude.json"
        data_dict = {
            "numStartups": 30,
            "includeCoAuthoredBy": False,
            "hasCompletedOnboarding": True,
            "lastOnboardingVersion": "0.2.37",
            "doctorShownAtSession": 26,
            "lastReleaseNotesSeen": "0.2.37",
            "bypassPermissionsModeAccepted": True,
            "mcpServers": {
                **additional_mcp_servers
            }
        }

        if use_anthropic_api_key:
            api_key = LLM_API_KEYS[LLM.CLAUDE]
            data_dict["primaryApiKey"] = api_key
            self.auth_mode = "api_key"
        else:
            self.auth_mode = "oauth"

        log.debug("Claude is using auth mode: ", self.auth_mode)

        home_dir = str(Path.home())
        claude_config_path = os.path.join(home_dir, filename)
        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)
                    data_dict["mcpServers"] = {
                        **data_dict["mcpServers"],
                        **mcp_servers
                    }

        if not os.path.exists(claude_config_path):
            with open(claude_config_path, 'w') as f:
                json.dump(data_dict, f, indent=4)

        log.debug("*"*100)
        log.debug("CLAUDE CONFIG")
        log.debug(json.dumps(data_dict, indent=4))
        log.debug("*"*100)
                
        return claude_config_path

    def query(self, query: str):
        self._start()
        is_interrupt_triggered = self.is_interrupt_triggered
        if is_interrupt_triggered:
            self.is_interrupt_triggered = False

        # Restart log thread if needed for new query
        self._restart_log_thread()

        should_resume = self.is_session_active

        # Create temp file and write query to it
        temp_file_path = tempfile.mktemp()
        with open(temp_file_path, 'w') as temp_file:
            temp_file.write(query)
            temp_file.flush()
        
        args = [
            "--verbose",
            "--dangerously-skip-permissions",
            "-p",
            "--output-format",
            "stream-json",
        ]

        if self.config.model:
            args.append("--model")
            args.append(self.config.model)

        if should_resume:
            log.debug("-"*100)
            log.debug("Resuming session")
            log.debug("-"*100)
            args.append("--continue")
            try:
                notification = NotifyResumeArgs()
                self.notify(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

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

        claude_command = f"cat {temp_file_path} | claude {' '.join(args)} >> {temp_file_log}"
        log.debug("*"*100)
        log.debug(f"Running claude command: {claude_command}")
        log.debug("*"*100)

        blocklist_names = []

        if self.auth_mode == "oauth":
            blocklist_names.append("ANTHROPIC_API_KEY")

        log.debug(f"Blocklist names: {blocklist_names}")

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

        def on_complete(cmd_output):

            stdout_output = cmd_output.stdout
            stderr_output = cmd_output.stderr
            pid = cmd_output.pid
            return_code = cmd_output.return_code

            if self.is_interrupt_triggered:
                return

            self.status = self.AgentStatus.TURNS_COMPLETED

            log.debug(f"Claude agent subprocess {pid} has exited with code {return_code}")
            log.debug(f"Claude agent subprocess err output: {stderr_output}")
            log.debug(f"Claude 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.warning("Warning: Log thread did not finish within timeout")

            if self.is_interrupt_triggered:
                return

            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(">>>Before notify complete!!!")

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

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

        out.register_callback(on_complete)
        
        self.pid = out.pid
        self.is_session_active = True

        return out
