import tempfile
import os
import json
import contextlib
from pathlib import Path
from blocks import bash
import time
import threading
from blocks_control_sdk.utils.logger import log, is_debug
from blocks_control_sdk.utils.main import is_bash_error_code, render_agent_error_message
from blocks_control_sdk.control.agent_base import CodingAgentBaseCLI, NotifyToolCallArgs, NotifyMessageArgs, NotifyCompleteArgs, NotifyStartArgs, NotifyResumeArgs
from blocks_control_sdk.parsers.gemini_messages import to_model_response
from blocks_control_sdk.parsers.base_messages import BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE, print_litellm_model_response

class GeminiCLI(CodingAgentBaseCLI):

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

        # 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

        # Message aggregation - simple buffer without locks (single-threaded processing)
        self._content_buffer = []

    def _flush_content_buffer(self):
        """Flush the content buffer and send as a single aggregated message."""
        if self._content_buffer:
            combined_content = "".join(self._content_buffer).strip()
            if combined_content:
                from litellm.types.utils import Message
                message = Message(role="assistant", content=combined_content)
                self.assistant_messages.append(message)
                self.notify(NotifyMessageArgs(message=message))
                log.debug(f"Flushed aggregated content ({len(self._content_buffer)} fragments): {combined_content[:100]}...")
            self._content_buffer = []

    def new_chat_thread(self, chat_thread_id: str = None):
        """Override to reset buffer state for new chat thread."""
        super().new_chat_thread(chat_thread_id)

        # Create a new temp log file for the new chat thread
        self.temp_file_log = tempfile.mktemp()
        self.log_file_position = 0

        # Clear content buffer for fresh start
        self._content_buffer = []

        # Stop and restart log thread with fresh state
        if self.log_thread and self.log_thread.is_alive():
            self.log_thread_should_stop.set()
            self.log_thread_finished.wait(timeout=2)

        log.debug(f"New chat thread initialized with fresh log file: {self.temp_file_log}")

    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 line: {e}")
            return

        # Try to convert to ModelResponse - this is what we care about for regression testing
        try:
            log.debug("*" * 100)
            model_response = to_model_response(raw_json)
            print_litellm_model_response(model_response)

            model_response.choices[0].message.provider_specific_fields.get("delta")

            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.role == "assistant" and message.content:
                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
            )

            # Handle tool calls - flush buffer before processing
            if tools:
                self._flush_content_buffer()

                for tool in tools:
                    log.debug(f"Tool: {tool.function.name}")
                    log.debug(f"Tool arguments: {tool.function.arguments}")
                    try:
                        parsed_args = json.loads(tool.function.arguments)
                        parsed_args = {
                            "__name__": tool.function.name,
                            **parsed_args,
                        }
                        notification = NotifyToolCallArgs(
                            tool_name=tool.function.name,
                            serialized_args=json.dumps(parsed_args),
                        )
                        log.debug(f"Notifying tool call v2: {notification}")
                        self.notify(notification)
                    except Exception as e:
                        log.error(f"Error notifying tool call v2: {e}")

            # Buffer assistant messages instead of sending immediately
            if (
                message.role == "assistant" and
                message.content and
                message.provider_specific_fields.get(BLOCKS_PROVIDER_SPECIFIC_FIELDS__IS_FINAL_MESSAGE, False) is False
            ):
                self._content_buffer.append(message.content)
            else:
                self._flush_content_buffer()

            log.debug("*" * 100)

        except ValueError as e:
            # Handle cases where JSON is valid but format is not supported by to_model_response
            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(f"Error converting to model response: {e}")
                log.error(f"Raw JSON: {line}")
        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(f"Error processing stream JSON line: {e}")

    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 write_config(self, additional_mcp_servers={}) -> str:
        """
        Write the gemini config file if it doesn't exist.
        """
        log.debug("="*100)
        log.debug("Writing Gemini config")
        log.debug("="*100)
        workspace_dir = str(self.workspace_dir.absolute())
        filename = ".gemini/settings.json"
        data_dict = {
            "theme": "GitHub",
            "selectedAuthType": "gemini-api-key",
            "includeDirectories": [
                "/tmp",
                "/bundle",
                workspace_dir
            ],
            "tools": {
                "exclude": ["save_memory"]
            },
            "mcpServers": {
                **additional_mcp_servers
            }
        }
        home_dir = str(Path.home())
        gemini_config_path = os.path.join(home_dir, filename)
        mcp_dir = os.path.join(home_dir, ".config/blocks/mcp.json")

        Path(Path().home() / ".gemini").mkdir(parents=True, exist_ok=True)
        
        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(gemini_config_path):
            with open(gemini_config_path, 'w') as f:
                json.dump(data_dict, f, indent=4)

        log.debug("="*100)
        log.debug("Gemini config written")
        log.debug(Path(gemini_config_path).read_text())
        log.debug("="*100)
                
        return gemini_config_path

    def interrupt(self):
        if self.pid is not None:
            self._kill_process(self.pid)
            self.pid = None
            self.is_interrupt_triggered = True
            # Clear stale buffered content on interrupt
            self._content_buffer = []

    def query(self, query: str):
        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 = [
            "-y",
            "-p",
            f'"$(cat {temp_file_path})"',
            "--output-format stream-json"
        ]

        if should_resume:
            log.debug("-"*100)
            log.debug("Resuming session")
            log.debug("-"*100)
            args.append("-r latest")
            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}")

        gemini_command = f"gemini {' '.join(args)} >> {temp_file_log}"
        log.debug("*"*100)
        log.debug(f"Running gemini command: {gemini_command}")
        log.debug("*"*100)

        out = bash(
            gemini_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"Gemini agent subprocess {pid} has exited with code {return_code}")
            log.debug(f"Gemini agent subprocess err output: {stderr_output}")
            log.debug(f"Gemini 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

            # Flush any remaining buffered content before completing
            self._flush_content_buffer()

            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)

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

        out.register_callback(on_complete)

        self.pid = out.pid
        self.is_session_active = True

        return out
      