import tempfile
from typing import Optional, Dict, Any, List
import os
import json
import contextlib
from pathlib import Path
from blocks import bash
import json
from blocks_control_sdk.utils import get_unset_command, is_bash_error_code, render_agent_error_message, to_pascal_case
import time
import threading
import shutil
import fcntl
from blocks_control_sdk.logger import log, stream_line
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__TODOWRITE,
    to_model_response,
    AnthropicTodoListUpdate
)
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, AnthropicAuthenticationMode

class ClaudeAgentConfig(BaseModel):
    model: AnthropicModels = AnthropicModels.sonnet
    authentication_mode: AnthropicAuthenticationMode = AnthropicAuthenticationMode.auto

class ClaudeCodeCLIExp(CodingAgentBaseCLI):

    has_oauth_credentials: bool = False

    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
        self.system_prompt = None

        # 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

    @property
    def resolved_authentication_mode(self) -> AnthropicAuthenticationMode:
        mode = self.config.authentication_mode
        if mode != AnthropicAuthenticationMode.auto:
            return mode
        # Auto precedence: bedrock > oauth > api_key
        if os.getenv("AWS_BEARER_TOKEN_BEDROCK"):
            return AnthropicAuthenticationMode.bedrock
        if self.has_oauth_credentials:
            return AnthropicAuthenticationMode.oauth
        return AnthropicAuthenticationMode.api_key

    def _process_log_line_for_messages(self, line):
        """Process a single log line for messages and notifications"""
        stream_line(line)
        try:
            log.debug("*"*100)
            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 []

            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_v2(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_v2(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:
                                _updated_name = to_pascal_case(tool.function.name.replace("mcp__blocks-internal-mcp__", ""))
                                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}")
                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_v2(notification)
                except Exception as e:
                    log.error(f"Error notifying message v2: {e}")

            log.debug("*"*100)
        except Exception as e:
            log.debug("*"*100)
            log.error(f"Error parsing line: {e}")
            log.debug("*"*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):
        super().interrupt()

    def _hot_reload_cleanup(self):
        """Reset log thread state for clean restart during hot reload."""
        # Stop log thread gracefully so it doesn't process stale data
        if self.log_thread and self.log_thread.is_alive():
            self.log_thread_should_stop.set()
            self.log_thread_finished.wait(timeout=2)

        # Reset interrupt flag for new query
        self.is_interrupt_triggered = False

    def write_config(self, additional_mcp_servers={}, has_oauth_credentials: bool = False) -> str:
        """
        Write the claude config file if it doesn't exist.
        """
        self.has_oauth_credentials = has_oauth_credentials

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

        resolved_mode = self.resolved_authentication_mode
        home_dir = str(Path.home())
        credentials_path = os.path.join(home_dir, ".claude", ".credentials.json")

        if resolved_mode == AnthropicAuthenticationMode.api_key:
            api_key = LLM_API_KEYS[LLM.CLAUDE]
            data_dict["primaryApiKey"] = api_key
            with contextlib.suppress(FileNotFoundError):
                os.remove(credentials_path)
        elif resolved_mode == AnthropicAuthenticationMode.bedrock:
            with contextlib.suppress(FileNotFoundError):
                os.remove(credentials_path)

        log.info(f"Claude is using authentication mode: {resolved_mode.value}")

        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.info(json.dumps(data_dict, indent=4))
        log.debug("*"*100)

        return claude_config_path

    # Hook event types supported by Claude Code
    HOOK_EVENTS = [
        "PreToolUse", "PostToolUse", "PermissionRequest",
        "Notification", "UserPromptSubmit", "Stop",
        "SubagentStop", "PreCompact", "SessionStart", "SessionEnd",
        "PostClone"  # Custom hook for post-clone scripts
    ]

    @staticmethod
    def _load_settings(folder_path: str) -> Dict[str, Any]:
        """Load settings.json from a .claude folder path."""
        if os.path.isfile(folder_path):
            settings_path = folder_path
        else:
            settings_path = os.path.join(folder_path, "settings.json")

        if not os.path.exists(settings_path):
            return {}

        with open(settings_path, 'r') as f:
            return json.load(f)

    @staticmethod
    def _merge_hook_arrays(
        target_hooks: List[Dict[str, Any]],
        source_hooks: List[Dict[str, Any]],
        source_project_root: str
    ) -> List[Dict[str, Any]]:
        """
        Merge two hook event arrays.
        Target hooks come first for same matchers (precedence via position).
        """
        # Build a map: matcher -> {"target_hooks": [], "source_hooks": []}
        matcher_map: Dict[str, Dict[str, List]] = {}

        # Process target hooks first (they get precedence)
        for item in target_hooks:
            if isinstance(item, dict):
                matcher = item.get("matcher", "*")
                if matcher not in matcher_map:
                    matcher_map[matcher] = {"target_hooks": [], "source_hooks": []}
                # Add the individual hooks from this matcher entry
                hooks_list = item.get("hooks", [])
                matcher_map[matcher]["target_hooks"].extend(hooks_list)

        # Process source hooks
        for item in source_hooks:
            if isinstance(item, dict):
                matcher = item.get("matcher", "*")
                if matcher not in matcher_map:
                    matcher_map[matcher] = {"target_hooks": [], "source_hooks": []}
                hooks_list = item.get("hooks", [])

                def transform_command_hook(hook: Dict[str, Any]) -> Dict[str, Any]:
                    _hook_command = hook.get("command")
                    final_hook_command = f"{_hook_command}"
                    return {
                        "type": "command",
                        "command": final_hook_command,
                    }
                # for hooks.type == "command", list them
                command_hooks = [transform_command_hook(hook) for hook in hooks_list if hook.get("type") == "command"]
                other_hooks = [hook for hook in hooks_list if hook.get("type") != "command"]
                matcher_map[matcher]["source_hooks"].extend(command_hooks + other_hooks)

        # Build result: target hooks first, then source hooks
        result = []
        for matcher, hooks_data in matcher_map.items():
            combined_hooks = hooks_data["target_hooks"] + hooks_data["source_hooks"]
            if combined_hooks:
                result.append({
                    "matcher": matcher,
                    "hooks": combined_hooks
                })

        return result

    @staticmethod
    def _merge_hooks(
        target_settings: Dict[str, Any],
        source_settings: Dict[str, Any],
        source_project_root: str
    ) -> Dict[str, Any]:
        """Merge hooks from two settings dictionaries. Target hooks come first."""
        target_hooks = target_settings.get("hooks", {})
        source_hooks = source_settings.get("hooks", {})

        merged_hooks = {}

        # Get all unique event types from both
        all_events = set(target_hooks.keys()) | set(source_hooks.keys())

        for event in all_events:
            if event in ClaudeCodeCLIExp.HOOK_EVENTS:
                target_arr = target_hooks.get(event, [])
                source_arr = source_hooks.get(event, [])
                merged_hooks[event] = ClaudeCodeCLIExp._merge_hook_arrays(target_arr, source_arr, source_project_root)

        return {"hooks": merged_hooks}

    @staticmethod
    def merge_claude_folder_into_target(source_folder: str, target_folder: str, source_project_root: str) -> bool:
        """
        Merge source .claude folder INTO existing target .claude folder.

        Target's existing config takes precedence:
        - Target hooks appear FIRST in arrays (so exit 2 hooks block)
        - Source files are copied but DON'T overwrite existing target files

        Args:
            source_folder: Path to .claude folder from cloned repo
            target_folder: Path to existing target .claude folder (has internal config)
            source_project_root: Path to the project root of the source folder

        Returns:
            True if merge succeeded, False if failed (target unchanged on failure)
        """
        # Ensure target directory exists for lock file
        os.makedirs(target_folder, exist_ok=True)

        lock_file_path = os.path.join(target_folder, ".merge.lock")

        with open(lock_file_path, 'w') as lock:
            try:
                # Try to acquire exclusive lock (non-blocking first)
                fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                # Another merge in progress, wait for it
                log.info("[Claude Merge] Merge lock held by another process, waiting...")
                fcntl.flock(lock.fileno(), fcntl.LOCK_EX)

            try:
                # Step 1: Load settings from both folders
                target_settings = ClaudeCodeCLIExp._load_settings(target_folder)
                source_settings = ClaudeCodeCLIExp._load_settings(source_folder)

                # Step 2: Copy source files to target (don't overwrite existing)
                # Use a recursive helper to merge directories properly
                def copy_without_overwrite(src_dir: str, dst_dir: str):
                    """Recursively copy files from src to dst, not overwriting existing files."""
                    os.makedirs(dst_dir, exist_ok=True)
                    for item in os.listdir(src_dir):
                        src_path = os.path.join(src_dir, item)
                        dst_path = os.path.join(dst_dir, item)
                        if os.path.isdir(src_path):
                            # Recursively merge directories
                            copy_without_overwrite(src_path, dst_path)
                        elif not os.path.exists(dst_path):
                            # Copy file only if it doesn't exist
                            shutil.copy2(src_path, dst_path)

                if os.path.exists(source_folder) and os.path.isdir(source_folder):
                    copy_without_overwrite(source_folder, target_folder)

                # Step 3: Merge hooks (target first for precedence)
                merged_settings = ClaudeCodeCLIExp._merge_hooks(target_settings, source_settings, source_project_root)

                # Step 4: Validate merged settings before writing
                def validate_settings(settings: Dict[str, Any]) -> bool:
                    """Validate that settings has the basic schema."""
                    if not isinstance(settings, dict):
                        return False
                    if "hooks" not in settings:
                        return False
                    if not isinstance(settings["hooks"], dict):
                        return False
                    return True

                # If merged settings are invalid, fall back to target settings
                if not validate_settings(merged_settings):
                    log.info(f"[Claude Merge] Warning: Merged settings validation failed, using target settings only")
                    merged_settings = {"hooks": target_settings.get("hooks", {})}

                # Step 5: Write settings.json to target
                target_settings_path = os.path.join(target_folder, "settings.json")
                with open(target_settings_path, 'w') as f:
                    json.dump(merged_settings, f, indent=2)
                    f.write('\n')

                log.info(f"[Claude Merge] Successfully merged source .claude folder into {target_folder}")
                return True

            except Exception as e:
                log.error(f"[Claude Merge] Error during merge: {e}. Target folder unchanged.")
                return False

            finally:
                # Release lock
                fcntl.flock(lock.fileno(), fcntl.LOCK_UN)

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

        # Restart log thread if needed for new query (only if tailing is enabled)
        if tail:
            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.system_prompt:
            args.append("--system-prompt-file")
            system_prompt_temp_file_path = tempfile.mktemp()
            with open(system_prompt_temp_file_path, 'w') as temp_file:
                temp_file.write(self.system_prompt)
                temp_file.flush()
            args.append(system_prompt_temp_file_path)

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

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


        resolved_mode = self.resolved_authentication_mode

        additional_env_vars = {}
        additional_blocklist_prefixes = []

        if resolved_mode == AnthropicAuthenticationMode.bedrock:
            AWS_BEARER_TOKEN_BEDROCK = os.getenv("AWS_BEARER_TOKEN_BEDROCK", None)
            if AWS_BEARER_TOKEN_BEDROCK:
                additional_env_vars["AWS_BEARER_TOKEN_BEDROCK"] = AWS_BEARER_TOKEN_BEDROCK
            additional_env_vars["CLAUDE_CODE_USE_BEDROCK"] = "1"

        if resolved_mode not in [AnthropicAuthenticationMode.api_key, AnthropicAuthenticationMode.auto]:
            additional_blocklist_prefixes.append("ANTHROPIC_API_KEY")

        out = bash(
            claude_command,
            background=self.is_background,
            blocklist_prefixes=["BLOCKS_BUNDLE_PRESIGNED_","AWS_", "ECS_", "E2B", *additional_blocklist_prefixes],
            set_env={
                "AWS_PROFILE": "blocks",
                **additional_env_vars
            }
        )

        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 return_code != 0:
            #     print("!"*100)
            #     print(f"WARNING!!! Retrying agent execution")
            #     print(f"Claude agent subprocess {pid} has exited with code {return_code}")
            #     print(f"Claude agent subprocess err output: {stderr_output}")
            #     print(f"Claude agent subprocess output: {stdout_output}")
            #     print("!"*100)
            #     self.query(query)
            #     return

            if self.is_interrupt_triggered:
                return

            self.status = self.AgentStatus.TURNS_COMPLETED

            log.info(f"Claude agent subprocess {pid} has exited with code {return_code}")
            log.info(f"Claude agent subprocess err output: {stderr_output}")
            log.info(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.info(f"Signaling log thread to stop... [assistant messages: {len(self.assistant_messages)}]")
            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.info(f"Log thread finished successfully [assistant messages: {len(self.assistant_messages)}]")
            else:
                log.info(f"Warning: Log thread did not finish within timeout [assistant messages: {len(self.assistant_messages)}]")

            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_v2(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
        self._set_background_command_output(out)

        return out
