import os
import tempfile
import json
import tomlkit
from typing import Union
import contextlib
from blocks import bash
from pathlib import Path
import threading
import time
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.parsers.openai_messages import to_model_response
from blocks_control_sdk.parsers.base_messages import print_litellm_model_response
from pydantic import BaseModel
from blocks_control_sdk.control.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


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

class Codex(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

        def tail_log_file():
            while True:

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

                            # Try to convert to ModelResponse - this is what we care about for regression testing
                            try:
                                log.debug("*"*100)
                                log.debug(f"Processing position#: {self.logfile_position}")
                                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 len(tools) > 0:
                                    parsed_type = "tool_call"
                                elif message.role == "tool":
                                    parsed_type = "tool_output"
                                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
                                )

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

                                    continue
                                elif message.role == "tool":
                                    log.debug(f"TOOL OUTPUT: {message.content}")
                                    continue
                                elif message.role == "assistant" and message.content:
                                    self.assistant_messages.append(message)
                                    try:
                                        notification = NotifyMessageArgs(
                                            message=message,
                                        )
                                        self.notify(notification)
                                    except Exception as e:
                                        log.error(f"Error notifying message: {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(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):
        if self.pid is not None:
            self._kill_process(self.pid)
            self.pid = None
            self.is_interrupt_triggered = True

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

        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 = [
            "--skip-git-repo-check",
            "--dangerously-bypass-approvals-and-sandbox",
        ]
        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.debug("-"*100)
            log.debug(f"Resuming session: {last_session_file_name}")
            log.debug("-"*100)
            try:
                notification = NotifyResumeArgs()
                self.notify(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(notification)
        except Exception as e:
            log.error(f"Error notifying start v2: {e}")

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

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

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

            last_assistant_message = self.assistant_messages[-1]
            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

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

        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.
        """
        root_dir = Path.home() / ".codex" / "sessions"
        root_dir.mkdir(parents=True, exist_ok=True)
        latest_file = None
        latest_file_name = None
        latest_mtime = -1
        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)
                        if mtime > latest_mtime:
                            latest_mtime = mtime
                            latest_file = full_path
                            latest_file_name = filename
                    except FileNotFoundError:
                        continue 

        return latest_file if not return_file_name else latest_file_name

    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"]
            
            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"
            },
            "model_providers": {
                "openai-chat": {
                    "base_url": "https://api.openai.com/v1",
                    "env_key": "OPENAI_API_KEY",
                    "wire_api": "chat",
                    "name": "OpenAI using Chat Completions",
                },
                "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
            }
        }
        
        # 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,
            **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.debug("*"*100)
        log.debug("CODEX AUTH CONFIG")
        log.debug(auth_file_path.read_text())
        log.debug("*"*100)
        
        log.debug("*"*100)
        log.debug("CODEX CONFIG")
        log.debug(codex_config_path.read_text())
        log.debug("*"*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)
