import os
import tempfile
import json
import tomlkit
from typing import Union, List
import contextlib
from blocks import bash
from pathlib import Path
import threading
import time
from blocks_control_sdk.utils import is_bash_error_code, render_agent_error_message, to_pascal_case
from blocks_control_sdk.parsers.openai_messages import to_model_response
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.openai import OpenAIModels
from blocks_control_sdk.parsers.openai_messages import OPENAI_TOOLS__SHELL
from blocks_control_sdk.logger import log, stream_line


class CodexAgentConfig(BaseModel):
    model: OpenAIModels = OpenAIModels.gpt_5

class CodexAgentCLI(CodingAgentBaseCLI):

    def __init__(self, chat_thread_id: str = None, config: CodexAgentConfig = CodexAgentConfig()):
        super().__init__(chat_thread_id)
        self.config = config
        self.logfile_position = 0
        self.logfile_current_file = None
        self._tracked_session_ids: List[str] = []
        self._session_ids_lock = threading.Lock()
        self._stdout_events: List[dict] = []  # Generic event capture from stdout
        self._stdout_events_lock = threading.Lock()

        def tail_log_file():
            while True:
                # Process captured stdout events
                with self._stdout_events_lock:
                    if self._stdout_events:
                        for event in self._stdout_events:
                            event_type = event.get("type")

                            log.debug(f"[stdout_handler] Processing event: {event_type}")

                            # Handle different event types
                            # Note: turn.failed removed - errors are handled in on_complete via stderr
                            # Future event types can be added here:
                            # if event_type == "some.other.event":
                            #     handle_other_event(event)
                            pass

                        # Clear processed events
                        self._stdout_events.clear()

                latest_session_file = self._find_latest_session_file()
                if not latest_session_file:
                    time.sleep(0.5)
                    continue

                if latest_session_file != self.logfile_current_file:
                    self.logfile_position = 0
                    self.logfile_current_file = latest_session_file
                try:
                    with open(latest_session_file, "r") as f:
                        f.seek(self.logfile_position)
                        for line in f:
                            stream_line(line)
                            try:
                                log.debug("*"*100)
                                log.debug(f"Processing position#: {self.logfile_position}")
                                print_litellm_model_response(
                                    to_model_response(
                                        json.loads(line)
                                    )
                                )
                                model_response = to_model_response(
                                    json.loads(line)
                                )
                                message = model_response.choices[0].message
                                tools = message.tool_calls or []

                                if len(tools) > 0:
                                    for tool in tools:
                                        if tool.function.name == OPENAI_TOOLS__SHELL:
                                            try:
                                                tool_args = tool.function.arguments
                                                parsed_tool_args = json.loads(tool_args)
                                                command = parsed_tool_args["command"]
                                                command_string = " ".join(command)

                                                self.notify_v2(NotifyToolCallArgs(
                                                    tool_name=tool.function.name,
                                                    serialized_args=json.dumps({
                                                        "__name__": tool.function.name,
                                                        "command": command_string,
                                                    }),
                                                ))
                                            except Exception as e:
                                                log.error(f"Error notifying tool call: {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:
                                                    _updated_name = to_pascal_case(tool.function.name.replace("mcp__blocks-internal-mcp__", ""))
                                                    # convert some_name_blah to SomeNameBlah
                                                    parsed_args = {
                                                        **parsed_args,
                                                        "__name__": _updated_name,
                                                    }

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

                                    continue
                                elif message.role == "tool":
                                    log.debug("TOOL OUTPUT:", message.content)
                                    continue
                                elif message.role == "assistant" and message.content:
                                    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}")

                                log.debug("*"*100)

                            except Exception as e:
                                log.error(f"Error parsing line: {e}")
                        self.logfile_position = f.tell()
                except FileNotFoundError:
                    # File might have been rotated/deleted
                    self.logfile_position = 0
                    self.logfile_current_file = None
                except Exception:
                    pass
                time.sleep(0.5)

        self.log_thread = threading.Thread(target=tail_log_file, daemon=True)
        self.log_thread.start()

    def _start(self):
        pass

    def interrupt(self):
        super().interrupt()

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

        NOTE: Do NOT clear _tracked_session_ids here - we need them to find
        the correct session file for resume. The daemon log thread continues
        running (no stop mechanism needed).
        """
        self.is_interrupt_triggered = False
        # Clear captured stdout events on hot reload
        with self._stdout_events_lock:
            self._stdout_events.clear()
        # Keep _tracked_session_ids intact for session continuity
        # Keep logfile_position and logfile_current_file for log continuity

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

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

        # 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 = [
            "--skip-git-repo-check",
            "--dangerously-bypass-approvals-and-sandbox",
            "--json"
        ]
        last_session_file_name = self._find_latest_session_file(return_file_name=True)
        # extract UUID from last_session_file_name using regex
        if should_resume and last_session_file_name:
            import re
            session_uuid = re.search(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', last_session_file_name).group(0)
            log.info("-"*100)
            log.info("Resuming session: ", last_session_file_name)
            log.info("-"*100)
            try:
                notification = NotifyResumeArgs()
                self.notify_v2(notification)
            except Exception as e:
                log.error(f"Error notifying resume v2: {e}")
            
            args.append(f'resume')
            args.append(session_uuid)
            args.append(f'"$(cat {temp_file_path})"')
        else:
            args.append(f'"$(cat {temp_file_path})"')


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

        log.info(f"Executing command: codex exec {' '.join(args)}")
        out = bash(
            f"codex exec {' '.join(args)}",
            background=True,
            blocklist_prefixes=["BLOCKS_BUNDLE_PRESIGNED_", "AWS_", "ECS_", "E2B"],
            set_env={
                "AWS_PROFILE": "blocks"
            }
        )

        # Capture thread_id from early stdout (thread.started comes first)
        def capture_thread_id():
            pos = None
            for _ in range(50):  # Try for ~5 seconds max
                result, pos = out.get_output_since_position(pos)
                if result.stdout:
                    for line in result.stdout.strip().split('\n'):
                        if not line:
                            continue
                        try:
                            event = json.loads(line)
                            if event.get("type") == "thread.started":
                                thread_id = event.get("thread_id")
                                if thread_id:
                                    with self._session_ids_lock:
                                        if thread_id not in self._tracked_session_ids:
                                            self._tracked_session_ids.append(thread_id)
                                    log.debug(f"Tracking session ID: {thread_id}")
                                    return  # Done, got what we needed
                        except json.JSONDecodeError:
                            pass
                time.sleep(0.1)

        # Only capture thread_id if tailing is enabled
        if tail:
            threading.Thread(target=capture_thread_id, daemon=True).start()

        # Capture generic JSONL events from stdout
        def capture_stdout_events():
            """Capture JSONL events from stdout stream for handling in log thread"""
            pos = None
            # Event types we want to capture (extensible list)
            # Note: turn.failed removed - errors are handled in on_complete via stderr
            capture_types = set()  # Reserved for future event types if needed

            while True:
                result, pos = out.get_output_since_position(pos)
                if result.stdout:
                    for line in result.stdout.strip().split('\n'):
                        if not line:
                            continue
                        try:
                            event = json.loads(line)
                            event_type = event.get("type")

                            # Capture events we care about
                            if event_type in capture_types:
                                with self._stdout_events_lock:
                                    self._stdout_events.append(event)
                                log.debug(f"[stdout_capture] Captured event: {event_type}")

                        except json.JSONDecodeError:
                            pass

                # Exit when process completes
                if out._completed:
                    break
                time.sleep(0.25)

        # Start the capture thread if tailing is enabled
        if tail:
            threading.Thread(target=capture_stdout_events, daemon=True).start()

        def on_complete(cmd_output):

            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"Codex agent subprocess {pid} has exited with code {return_code}")
            log.info(f"Codex agent subprocess err output: {stderr_output}")
            log.info(f"Codex agent subprocess output: {stdout_output}")

            # Grace period to allow the agent to send its final message
            time.sleep(7)

            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"[Codex agent] Last assistant message content (is bash error code: {is_bash_error_code(return_code)}): {last_assistant_message_content}")

            if is_bash_error_code(return_code):
                # Try to extract nice error message from turn.failed event in stdout
                turn_failed_error = None
                if stdout_output:
                    # only load the last 10 lines of stdout_output (safely)
                    stdout_output_lines = stdout_output.strip().split('\n')
                    stdout_output_lines = stdout_output_lines[-10:]
                    for line in stdout_output_lines:
                        if not line:
                            continue
                        try:
                            event = json.loads(line)
                            if event.get("type") == "turn.failed":
                                error_obj = event.get("error", {})
                                turn_failed_error = error_obj.get("message")
                                log.info(f"[Codex agent] Found turn.failed error: {turn_failed_error}")
                                break
                        except json.JSONDecodeError:
                            pass

                # Use turn.failed error if available, otherwise fall back to stderr
                last_assistant_message_content = turn_failed_error or last_assistant_message_content
                if "unexpected status 401 Unauthorized" in last_assistant_message_content:
                    last_assistant_message_content = "Error: There was an authentication error. Please double check your API Key valid or Subscription is active"
                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_v2(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
        self._set_background_command_output(out)

        return out

    def _find_latest_session_file(self, return_file_name=False):
        """
        Find the latest session file in the Codex sessions directory.

        Prioritizes session files matching tracked session IDs. Falls back to
        any session file if no tracked sessions are found.

        Note:
            - ~/.codex/history.jsonl also has the combined history of all sessions.
            - .codex/sessions/2025/07/27/rollout-2025-07-27T14-49-03-a0276d12-877a-4fb0-97ea-e11ad8be1683.jsonl Example file name.
        """
        import re
        root_dir = Path.home() / ".codex" / "sessions"
        root_dir.mkdir(parents=True, exist_ok=True)

        # Get tracked session IDs (thread-safe)
        with self._session_ids_lock:
            tracked_ids = self._tracked_session_ids.copy()

        uuid_pattern = re.compile(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}')

        all_sessions = []
        tracked_sessions = []

        for dirpath, _, filenames in os.walk(root_dir):
            for filename in filenames:
                if filename.endswith('.jsonl'):
                    full_path = os.path.join(dirpath, filename)
                    try:
                        mtime = os.path.getmtime(full_path)
                        match = uuid_pattern.search(filename)
                        file_uuid = match.group(0) if match else None

                        session_info = (full_path, filename, mtime)
                        all_sessions.append(session_info)

                        if file_uuid and file_uuid in tracked_ids:
                            tracked_sessions.append(session_info)
                    except FileNotFoundError:
                        continue

        # Prefer tracked sessions, fall back to all sessions
        candidates = tracked_sessions if tracked_sessions else all_sessions

        if not candidates:
            return None

        # Find latest by modification time
        latest = max(candidates, key=lambda x: x[2])
        return latest[1] if return_file_name else latest[0]

    def _json_mcp_to_toml_document(self, json_data: Union[dict, str]):
        """
        Convert JSON MCP server configuration to TOML format.
        
        Input format:
        {
            "server-name": {
                "command": "npx",
                "args": ["-y", "package-name"],
                "disabled": false,
                "autoApprove": [],
                "env": {"KEY": "value"}
            }
        }
        
        Output format:
        [mcp_servers.server-name]
        command = "npx"
        args = ["-y", "package-name"]
        env = { "KEY" = "value" }
        """
        if isinstance(json_data, str):
            data = json.loads(json_data)
        else:
            data = json_data
        
        doc = tomlkit.document()
        mcp_servers = tomlkit.table()
        
        for server_name, server_config in data.items():
            server_section = tomlkit.table()

            if "command" in server_config:
                server_section["command"] = server_config["command"]
            
            if "args" in server_config:
                server_section["args"] = server_config["args"]
            
            if "env" in server_config:
                env_inline = tomlkit.inline_table()
                for key, value in server_config["env"].items():
                    env_inline[key] = value
                server_section["env"] = env_inline
            
            if "disabled" in server_config:
                server_section["disabled"] = server_config["disabled"]
            
            if "autoApprove" in server_config:
                server_section["autoApprove"] = server_config["autoApprove"]

            server_section["tool_timeout_sec"] = 60*45 # 45 minutes
            
            mcp_servers[server_name] = server_section
        
        doc["mcp_servers"] = mcp_servers
        
        return doc

    def _build_config(self, additional_mcp_servers={}) -> dict:
        """
        Build the master config dictionary that exactly corresponds to the TOML structure.

        Returns:
            dict: Master configuration that matches TOML structure when serialized
        """
        blocks_mcp_config_path = Path.home() / ".config" / "blocks" / "mcp.json"

#

        # Base configuration (flattened to match TOML structure)
        master_config = {
            "model": self.config.model,
            # "model_provider": "openai-chat",
            "file_opener": "none",
            "sandbox_mode": "danger-full-access",
            "approval_policy": "never",
            "history": {
                "persistence": "save-all"
            },
            "features": {
                "web_search_request": True
            },
            "model_providers": {
                "openai-responses": {
                    "base_url": "https://api.openai.com/v1",
                    "env_key": "OPENAI_API_KEY",
                    "wire_api": "responses",
                    "name": "OpenAI using Chat Responses",
                },
            },
            "shell_environment_policy": {
                "inherit": "all",
                "ignore_default_excludes": True
            },
            "profiles": {
                "readonly": {
                    "sandbox_mode": "read-only",
                }
            },
        }

        # Load and merge MCP configuration
        mcp_servers = {}
        try:
            if blocks_mcp_config_path.exists():
                with open(blocks_mcp_config_path, "r") as f:
                    mcp_servers = json.load(f)
        except Exception as e:
            log.error(f"Error loading Blocks MCP config: {e}")

        # Merge with additional MCP servers
        mcp_servers = {
            **mcp_servers,
            **self._mcp_servers,
            **additional_mcp_servers
        }

        # Add MCP servers to master config if any exist
        if mcp_servers:
            master_config["mcp_servers"] = mcp_servers

        return master_config

    def write_config(self, additional_mcp_servers={}) -> str:
        """
        Write the config to both JSON and TOML files.
        
        Args:
            additional_mcp_servers: Additional MCP servers to merge with existing config
            
        Returns:
            str: Path to the generated TOML config file
        """
        codex_config_path = Path.home() / ".codex" / "config.toml"
        master_json_path = Path.home() / ".codex" / "master_config.json"
        auth_file_path = Path.home() / ".codex" / "auth.json"
        
        # Ensure directories exist
        codex_config_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Build master configuration
        master_config = self._build_config(additional_mcp_servers)
        
        # Save master config to JSON
        try:
            with open(master_json_path, "w") as f:
                json.dump(master_config, f, indent=2)
        except Exception as e:
            log.error(f"Error saving master config to JSON: {e}")

        # Save auth config to JSON
        try:
            # Check if existing auth.json has OAuth tokens
            existing_auth = {}
            if auth_file_path.exists():
                try:
                    with open(auth_file_path, "r") as f:
                        existing_auth = json.load(f)
                except Exception as e:
                    log.error(f"Error loading existing auth config: {e}")

            if "tokens" in existing_auth and existing_auth["tokens"]:
                # no-nop, auth config file already in place
                pass
            else:
                # No tokens exist, just write OPENAI_API_KEY
                with open(auth_file_path, "w") as f:
                    json.dump({
                        "OPENAI_API_KEY": os.getenv("OPENAI_API_KEY") or ""
                    }, f, indent=2)
        except Exception as e:
            log.error(f"Error saving auth config to JSON: {e}")
        
        # Convert to TOML using hybrid approach
        toml_content = self._convert_master_config_to_toml(master_config)
        
        # Write TOML config
        codex_config_path.write_text(toml_content)

        log.info("*"*100)
        log.info("CODEX AUTH CONFIG")
        log.info(auth_file_path.read_text())
        log.info("*"*100)

        log.info("*"*100)
        log.info("CODEX CONFIG")
        log.info(codex_config_path.read_text())
        log.info("*"*100)
        
        return str(codex_config_path)

    def patch_config(self, key: str, value):
        """
        Update a specific configuration value and rewrite both config files.
        
        Args:
            key: Configuration key (supports nested keys like 'codex_config.sandbox_mode')
            value: New value to set
        """
        master_json_path = Path.home() / ".codex" / "master_config.json"
        
        # Load existing master config
        master_config = {}
        try:
            if master_json_path.exists():
                with open(master_json_path, "r") as f:
                    master_config = json.load(f)
            else:
                # If no existing config, build a new one
                master_config = self._build_config()
        except Exception as e:
            log.error(f"Error loading existing config: {e}")
            master_config = self._build_config()
        
        # Update the specified key
        keys = key.split('.')
        current = master_config
        
        # Navigate to the parent of the target key
        for k in keys[:-1]:
            if k not in current:
                current[k] = {}
            current = current[k]
        
        # Set the final value
        current[keys[-1]] = value
        
        # Save updated master config to JSON
        try:
            with open(master_json_path, "w") as f:
                json.dump(master_config, f, indent=2)
        except Exception as e:
            log.error(f"Error saving patched config: {e}")
        
        # Regenerate TOML config
        self._write_toml_from_master_config(master_config)
    
    def _convert_master_config_to_toml(self, master_config: dict) -> str:
        """
        Convert master config to TOML using hybrid approach:
        - Regular sections: use tomlkit.dumps
        - MCP servers: use special inline table handling
        """
        toml_content = ""
        
        # Extract MCP servers if present
        mcp_servers = master_config.get("mcp_servers", {})
        
        # Create config without MCP servers for regular conversion
        config_without_mcp = {k: v for k, v in master_config.items() if k != "mcp_servers"}
        
        # Convert non-MCP sections using regular tomlkit
        if config_without_mcp:
            toml_content += tomlkit.dumps(config_without_mcp)
        
        # Convert MCP servers using special logic for inline tables
        if mcp_servers:
            mcp_toml_doc = self._json_mcp_to_toml_document(mcp_servers)
            if toml_content:
                toml_content += "\n"
            toml_content += tomlkit.dumps(mcp_toml_doc)
        
        return toml_content

    def _write_toml_from_master_config(self, master_config: dict):
        """
        Helper method to write TOML config from master config dict.
        """
        codex_config_path = Path.home() / ".codex" / "config.toml"
        
        # Use hybrid conversion approach
        toml_content = self._convert_master_config_to_toml(master_config)
        
        # Write TOML config
        codex_config_path.write_text(toml_content)
