"""RPC interface for Dash AI Chat in dev tools."""

import importlib
import json
import os
import time
from typing import Any, Dict, List, Optional, cast

from typing_extensions import Literal, NotRequired, TypedDict

from ._amplitude import ChatEvent, amplitude_service
from ._chat_agent import AGENT_TOOLS, ChatAgent, StagedChanges
from ._cloud_env import cloud_config
from ._llm_client import ChatMessage, LLMClient
from ._oauth import OAuthClient
from .exceptions import EntitlementError, TokenError

# Entitlement flag ID for Dash AI Chat access
DASH_AI_EXTENSION_ENTITLEMENT = "dash-ai-access"

# State file name for persisting chat state between hot-reloads
CHAT_STATE_FILE = ".dash-chat-state.json"


class ChatOperation(TypedDict):
    """RPC operation structure for chat."""

    operation: Literal[
        "initialize",
        "send_message",
        "get_staged_changes",
        "apply_changes",
        "clear_changes",
        "revert_change",
        "add_frontend_log",
        "authenticate",
        "auth_poll",
        "logout",
    ]
    data: Any


class ChatRPCResponse(TypedDict):
    """Standard RPC response structure."""

    result: NotRequired[Any]
    error: NotRequired[str]


class InitializeResult(TypedDict):
    """Result from initialize operation."""

    authenticated: bool
    project_path: str


class SendMessageResult(TypedDict):
    """Result from send_message operation."""

    assistant_message: str
    staged_changes: StagedChanges
    tool_calls_made: List[str]


class ApplyResult(TypedDict):
    """Result from apply_changes operation."""

    applied: List[Dict[str, Any]]


class DashAIChatRPC:
    """RPC handler for Dash AI Chat operations in dev tools.

    This handler manages the chat conversation, LLM interactions, and tool
    execution for the Dash AI Chat feature in dev tools.
    """

    # Maximum tool call iterations to prevent runaway agents
    MAX_ITERATIONS = 10

    def __init__(self) -> None:
        """Initialize the RPC handler."""
        client_id = cloud_config.get_oauth_client_id()
        self.oauth_client = OAuthClient(client_id)
        self._app_setup: Any = None
        self._agent: Optional[ChatAgent] = None
        self._conversation_history: List[ChatMessage] = []
        self._project_path: str = ""

    def _get_state_file_path(self) -> str:
        """Get the path to the state file."""
        return os.path.join(self._project_path, CHAT_STATE_FILE)

    def _save_state(self) -> None:
        """Save current chat state to file."""
        if not self._project_path:
            return

        state = {
            "conversation_history": self._conversation_history,
            "staged_changes": self._agent.get_staged_changes() if self._agent else {"edits": []},
        }

        try:
            state_path = self._get_state_file_path()
            with open(state_path, "w", encoding="utf-8") as f:
                json.dump(state, f, indent=2)
        except Exception:
            # Silently fail - state persistence is not critical
            pass

    def _load_state(self) -> bool:
        """Load chat state from file.

        Returns:
            True if state was loaded successfully, False otherwise
        """
        if not self._project_path:
            return False

        state_path = self._get_state_file_path()
        if not os.path.exists(state_path):
            return False

        try:
            with open(state_path, encoding="utf-8") as f:
                state = json.load(f)

            # Restore conversation history
            if "conversation_history" in state and state["conversation_history"]:
                self._conversation_history = state["conversation_history"]
                # Ensure system prompt has cache_control (may be missing from old states)
                if self._conversation_history and self._conversation_history[0].get("role") == "system":
                    self._conversation_history[0]["cache_control"] = {"type": "ephemeral"}

            # Restore staged changes to agent
            if self._agent and "staged_changes" in state:
                self._agent.staged_changes = state["staged_changes"]

            return True
        except Exception:
            # If loading fails, start fresh
            return False

    def _clear_state(self) -> None:
        """Delete the state file."""
        if not self._project_path:
            return

        state_path = self._get_state_file_path()
        try:
            if os.path.exists(state_path):
                os.remove(state_path)
        except Exception:
            pass

    def _get_project_path(self) -> str:
        """Get the project path from the running Dash app.

        Returns:
            Absolute path to the project directory
        """
        import dash

        app = None
        try:
            app = dash.get_app()
        except Exception:
            app = self._app_setup

        if app:
            try:
                app_module = importlib.import_module(app.config.name)
                return os.path.dirname(str(app_module.__file__))
            except Exception:
                pass

        return os.getcwd()

    def _get_project_structure(self) -> str:
        """Get a listing of project files for the system prompt.

        Returns:
            String listing of project files
        """
        if not self._project_path:
            return ""

        # Directories to skip
        skip_dirs = {
            "node_modules", ".git", "__pycache__", ".pytest_cache",
            "venv", ".venv", "env", ".env", "dist", "build", ".next",
            ".nuxt", "coverage", ".tox", "eggs", "*.egg-info",
        }

        files = []
        try:
            for root, dirs, filenames in os.walk(self._project_path):
                # Filter out directories to skip
                dirs[:] = [d for d in dirs if d not in skip_dirs and not d.startswith('.')]

                rel_root = os.path.relpath(root, self._project_path)
                for fname in filenames:
                    if fname.startswith('.'):
                        continue
                    if rel_root == '.':
                        files.append(fname)
                    else:
                        files.append(os.path.join(rel_root, fname))

                # Limit to prevent huge listings
                if len(files) > 100:
                    files.append("... (truncated)")
                    break
        except Exception:
            return ""

        return "\n".join(sorted(files))

    def _get_system_prompt(self) -> str:
        """Get the system prompt for the chat agent.

        Returns:
            System prompt string
        """
        project_structure = self._get_project_structure()

        return f"""You are Dash Chat, an AI assistant integrated into the Dash development environment.
You help developers build and debug their Dash applications.

Project path: {self._project_path}

Project structure:
{project_structure}

You have access to the following tools:
- read_file: Read contents of project files (large files show only relevant sections)
- stage_file_edit: Stage targeted edits using find-and-replace (old_content -> new_content)
- apply_staged_changes: DO NOT USE - the user has an Apply button in the UI
- list_files: List files in directories (use only if you need to explore subdirectories not shown above)
- run_command: Execute shell commands
- get_logs: View backend/frontend logs from the running Dash app
- search_files: Search for patterns across project files

When making file edits:
1. Read the relevant portion of the file first
2. Use stage_file_edit with the EXACT text to find (old_content) and replacement (new_content)
3. Only include the specific lines being changed, not the entire file
4. The old_content must match exactly including whitespace and indentation
5. After staging changes, STOP and let the user review them in the UI

Example stage_file_edit call:
- old_content: "def foo():\\n    return 1"
- new_content: "def foo():\\n    return 2"

=== CRITICAL: DO NOT AUTO-APPLY CHANGES ===
NEVER call apply_staged_changes. The user has an "Apply" button in the UI.
After you stage changes with stage_file_edit, your job is DONE for that task.
Do NOT offer to apply them. Do NOT ask if you should apply them. Just explain what you staged and wait.
The user will click the Apply button themselves when they're ready.
If the user explicitly says "apply" or "save the changes", only then may you call apply_staged_changes.

CRITICAL: If you say you will do something (like "I'll read the file" or "Let me make that change" or "Now I need to..."), you MUST include the tool call in the same response. Never describe an intended action without actually performing it. Either do the action immediately OR ask for confirmation first - never announce an action and then stop.

Be helpful, concise, and focused on Dash development best practices.
Keep your responses focused and actionable."""

    async def _check_response_incomplete(
        self, llm_client: LLMClient, assistant_response: str
    ) -> bool:
        """Check if an assistant response indicates intent to act without acting.

        Uses a fast/small model to determine if the response says the assistant
        will do something but hasn't actually done it (no tool calls made).

        Args:
            llm_client: The LLM client to use for the check
            assistant_response: The assistant's response text to analyze

        Returns:
            True if the response seems incomplete (said it would act but didn't)
        """
        check_messages: List[ChatMessage] = [
            {
                "role": "system",
                "content": (
                    "You are a classifier. Analyze the assistant response and determine "
                    "if it indicates the assistant intends to take an action (like reading "
                    "a file, making an edit, running a command) but hasn't actually done it yet. "
                    "\n\n"
                    "Respond 'YES' if the response ends with stated intent to act but no action taken. "
                    "Examples that should return YES:\n"
                    "- 'I'll read the file now' (says will read but didn't)\n"
                    "- 'Let me make that change' (says will change but didn't)\n"
                    "- 'Now I need to update the CSS:' (announces next step but stopped)\n"
                    "- 'I'll update the configuration:' (colon indicates more was expected)\n"
                    "- Ending with a colon after describing an intended action\n\n"
                    "Respond 'NO' if:\n"
                    "- The assistant has already completed actions (staged changes, read files, etc.)\n"
                    "- The assistant is asking the user a question or waiting for confirmation\n"
                    "- The assistant says 'Would you like me to...' or 'Should I...'\n"
                    "- The assistant has staged changes and is explaining them\n"
                    "- The response is a complete answer, explanation, or summary\n"
                    "- The response describes what was done (past tense), not what will be done\n\n"
                    "Staged changes are a valid stopping point - the user needs to review and apply them."
                ),
            },
            {
                "role": "user",
                "content": f"Assistant response to analyze:\n\n{assistant_response}",
            },
        ]

        try:
            response = await llm_client.chat_completion(
                messages=check_messages,
                tools=None,
                model="claude-haiku-4-5",
                max_tokens=10,
                temperature=0.0,
            )
            content = response["choices"][0]["message"].get("content") or ""
            result = content.strip().upper()
            return result.startswith("YES")
        except Exception as e:
            # Log the error type for debugging, but assume response is complete
            # to avoid blocking the user
            import sys
            print(f"[Dash Chat] Incomplete check failed: {type(e).__name__}: {e}", file=sys.stderr)
            return False

    async def handle_operation(self, operation: ChatOperation) -> ChatRPCResponse:
        """Handle a chat operation from dev tools.

        Args:
            operation: The operation to perform with its data

        Returns:
            ChatRPCResponse with result or error
        """
        operation_name = operation["operation"]
        data = operation.get("data")

        if not hasattr(self, operation_name):
            return {"error": f"Unsupported operation: {operation_name}"}

        try:
            method = getattr(self, operation_name)
            return await method(data)
        except Exception as e:
            error_type = type(e).__name__
            return {"error": f"{error_type}: {str(e)}"}

    async def initialize(self, data: Any) -> ChatRPCResponse:
        """Initialize chat session.

        Args:
            data: Not used

        Returns:
            Response with authentication status, project path, and restored state
        """
        is_authenticated = await self.oauth_client.is_authenticated()

        try:
            await self.oauth_client.refresh_access_token()
        except TokenError:
            is_authenticated = False
            self.oauth_client.clear_credentials()

        self._project_path = self._get_project_path()

        # Check entitlement if authenticated
        has_entitlement = False
        if is_authenticated:
            try:
                api_base_url = cloud_config.get_api_base_url()
                has_entitlement = await self.oauth_client.check_entitlement(
                    DASH_AI_EXTENSION_ENTITLEMENT, api_base_url
                )
            except (TokenError, EntitlementError):
                # If entitlement check fails, treat as no entitlement
                has_entitlement = False

        # Initialize agent only if authenticated AND has entitlement
        if is_authenticated and has_entitlement:
            self._agent = ChatAgent(self._project_path)

        # Try to load existing state from file
        state_loaded = self._load_state()

        # If no state was loaded, initialize with system prompt
        if not state_loaded:
            self._conversation_history = [
                {
                    "role": "system",
                    "content": self._get_system_prompt(),
                    "cache_control": {"type": "ephemeral"},
                }
            ]

        # Get messages to restore (excluding system prompt and auto-continue messages)
        restored_messages = []
        restored_staged_changes: StagedChanges = {"edits": []}
        if state_loaded:
            # Track tool calls and last assistant message for current turn
            current_tool_calls: List[str] = []
            last_assistant_with_tools: Optional[Dict[str, Any]] = None
            showed_assistant_this_turn = False

            for msg in self._conversation_history:
                role = msg.get("role")
                content = msg.get("content", "")

                # Skip system messages and auto-continue
                if role == "system" or content.startswith("[AUTO_CONTINUE]"):
                    continue

                # Add user messages
                if role == "user":
                    # If we had tool calls but no final response, show last assistant msg
                    if current_tool_calls and not showed_assistant_this_turn and last_assistant_with_tools:
                        last_assistant_with_tools["tool_calls_made"] = current_tool_calls
                        restored_messages.append(last_assistant_with_tools)

                    restored_msg: Dict[str, Any] = {"role": "user", "content": content}
                    if "start_time" in msg:
                        restored_msg["start_time"] = msg["start_time"]
                    if "end_time" in msg:
                        restored_msg["end_time"] = msg["end_time"]
                    restored_messages.append(restored_msg)
                    # Reset for new turn
                    current_tool_calls = []
                    last_assistant_with_tools = None
                    showed_assistant_this_turn = False

                # Handle assistant messages
                if role == "assistant":
                    # Collect tool names if this message has tool_calls
                    if "tool_calls" in msg:
                        for tc in msg["tool_calls"]:
                            tool_name = tc.get("function", {}).get("name", "")
                            if tool_name:
                                current_tool_calls.append(tool_name)
                        # Track last assistant message with tool_calls (as fallback)
                        if content.strip():
                            last_assistant_with_tools = {"role": "assistant", "content": content}
                            if "start_time" in msg:
                                last_assistant_with_tools["start_time"] = msg["start_time"]
                            if "end_time" in msg:
                                last_assistant_with_tools["end_time"] = msg["end_time"]

                    # Show assistant messages WITHOUT tool_calls (these are final responses)
                    if "tool_calls" not in msg and content.strip():
                        restored_msg = {"role": "assistant", "content": content}
                        if current_tool_calls:
                            restored_msg["tool_calls_made"] = current_tool_calls
                            current_tool_calls = []
                        if "start_time" in msg:
                            restored_msg["start_time"] = msg["start_time"]
                        if "end_time" in msg:
                            restored_msg["end_time"] = msg["end_time"]
                        restored_messages.append(restored_msg)
                        showed_assistant_this_turn = True
                        last_assistant_with_tools = None

            # Handle final turn if it only has tool_calls
            if current_tool_calls and not showed_assistant_this_turn and last_assistant_with_tools:
                last_assistant_with_tools["tool_calls_made"] = current_tool_calls
                restored_messages.append(last_assistant_with_tools)

            if self._agent:
                restored_staged_changes = self._agent.get_staged_changes()

        return {
            "result": {
                "authenticated": is_authenticated,
                "has_entitlement": has_entitlement,
                "project_path": self._project_path,
                "state_restored": state_loaded,
                "restored_messages": restored_messages,
                "restored_staged_changes": restored_staged_changes,
            }
        }

    async def send_message(self, data: Any) -> ChatRPCResponse:
        """Process a user message and get agent response.

        Args:
            data: Dict with "message" key containing user message

        Returns:
            Response with assistant message, staged changes, and tool calls made
        """
        if not self._agent:
            return {"error": "Chat not initialized. Please authenticate first."}

        user_message = data.get("message", "") if data else ""
        if not user_message:
            return {"error": "Message is required"}

        # Add user message to history with timestamp
        user_timestamp = time.time()
        user_msg: Dict[str, Any] = {
            "role": "user",
            "content": user_message,
            "start_time": user_timestamp,
            "end_time": user_timestamp,
        }
        # Add cache breakpoint on first user message (caches system prompt + first message)
        if len(self._conversation_history) == 1:  # Only system prompt exists
            user_msg["cache_control"] = {"type": "ephemeral"}
        self._conversation_history.append(user_msg)

        tool_calls_made: List[str] = []
        assistant_content = ""

        llm_client = LLMClient(self.oauth_client)

        async with llm_client:
            # Set context for intelligent file filtering
            self._agent.set_context(llm_client, user_message)

            # Track when assistant processing starts
            assistant_start_time = time.time()

            # Initial LLM call
            response = await llm_client.chat_completion(
                messages=self._conversation_history,
                tools=AGENT_TOOLS,
            )

            # Handle tool calls iteratively
            iteration = 0

            while iteration < self.MAX_ITERATIONS:
                choice = response["choices"][0]
                message = choice["message"]
                finish_reason = choice["finish_reason"]

                if finish_reason == "tool_calls" and "tool_calls" in message:
                    # Add assistant message with tool calls to history
                    tool_call_time = time.time()
                    self._conversation_history.append(
                        {
                            "role": "assistant",
                            "content": message.get("content") or "",
                            "tool_calls": cast(List[Dict[str, Any]], message["tool_calls"]),
                            "start_time": assistant_start_time,
                            "end_time": tool_call_time,
                        }
                    )

                    # Execute each tool call
                    for tool_call in message["tool_calls"]:
                        tool_name = tool_call["function"]["name"]
                        tool_calls_made.append(tool_name)
                        tool_exec_start = time.time()
                        result = await self._agent.execute_tool(tool_call)
                        tool_exec_end = time.time()

                        # Add tool result to history
                        self._conversation_history.append(
                            {
                                "role": "tool",
                                "tool_call_id": tool_call["id"],
                                "content": result,
                                "start_time": tool_exec_start,
                                "end_time": tool_exec_end,
                            }
                        )

                    # Save state after tool calls to prevent data loss on failure
                    self._save_state()

                    # Reset start time for next assistant response
                    assistant_start_time = time.time()

                    # Get next response
                    response = await llm_client.chat_completion(
                        messages=self._conversation_history,
                        tools=AGENT_TOOLS,
                    )
                    iteration += 1
                else:
                    # No more tool calls from LLM - exit the loop
                    break

            # Process final response (after tool-calling loop)
            choice = response["choices"][0]
            message = choice["message"]
            assistant_content = message.get("content") or ""

            # Handle empty responses - use last assistant content or generate summary
            if not assistant_content.strip():
                # Find the last assistant message with content
                for hist_msg in reversed(self._conversation_history):
                    if hist_msg.get("role") == "assistant" and hist_msg.get("content", "").strip():
                        assistant_content = hist_msg["content"]
                        break

                # If still empty and we made tool calls, generate a summary
                if not assistant_content.strip() and tool_calls_made:
                    assistant_content = f"Done. Used tools: {', '.join(tool_calls_made)}"

                # If truly nothing to say, return error
                if not assistant_content.strip():
                    self._save_state()
                    return {
                        "error": "The assistant returned an empty response. Please try again."
                    }

            # Check if response indicates intent to act but no action taken
            # Only check once on the final response (not during tool-call iterations)
            if iteration < self.MAX_ITERATIONS - 1:
                seems_incomplete = await self._check_response_incomplete(
                    llm_client, assistant_content
                )

                if seems_incomplete:
                    # Add assistant message and prompt to continue
                    incomplete_end_time = time.time()
                    self._conversation_history.append(
                        {
                            "role": "assistant",
                            "content": assistant_content,
                            "start_time": assistant_start_time,
                            "end_time": incomplete_end_time,
                        }
                    )
                    # Add a follow-up to prompt continuation (marked as auto-continue)
                    auto_continue_time = time.time()
                    self._conversation_history.append(
                        {
                            "role": "user",
                            "content": "[AUTO_CONTINUE] Please proceed with the action you described.",
                            "start_time": auto_continue_time,
                            "end_time": auto_continue_time,
                        }
                    )
                    # Reset start time for final response
                    assistant_start_time = time.time()
                    # Get one more response
                    response = await llm_client.chat_completion(
                        messages=self._conversation_history,
                        tools=AGENT_TOOLS,
                    )
                    assistant_content = response["choices"][0]["message"].get("content") or ""

            # Save final response
            final_end_time = time.time()
            self._conversation_history.append(
                {
                    "role": "assistant",
                    "content": assistant_content,
                    "start_time": assistant_start_time,
                    "end_time": final_end_time,
                }
            )

        # Save state after processing
        self._save_state()

        return {
            "result": {
                "assistant_message": assistant_content,
                "staged_changes": self._agent.get_staged_changes(),
                "tool_calls_made": tool_calls_made,
                "user_start_time": user_timestamp,
                "user_end_time": user_timestamp,
                "assistant_start_time": assistant_start_time,
                "assistant_end_time": final_end_time,
            }
        }

    async def send_message_stream(self, data: Any):
        """Process a user message and stream responses as SSE events.

        Yields the same message format as stored in conversation_history,
        allowing identical rendering for live streaming and restored state.

        Args:
            data: Dict with "message" key containing user message

        Yields:
            Dict messages in conversation_history format
        """
        if not self._agent:
            yield {"event": "error", "data": {"error": "Chat not initialized. Please authenticate first."}}
            return

        user_message = data.get("message", "") if data else ""
        if not user_message:
            yield {"event": "error", "data": {"error": "Message is required"}}
            return

        # Add user message to history with timestamp
        user_timestamp = time.time()
        user_msg: Dict[str, Any] = {
            "role": "user",
            "content": user_message,
            "start_time": user_timestamp,
            "end_time": user_timestamp,
        }
        # Add cache breakpoint on first user message
        if len(self._conversation_history) == 1:
            user_msg["cache_control"] = {"type": "ephemeral"}
        self._conversation_history.append(user_msg)

        # Yield user message
        yield {"event": "message", "data": user_msg}

        tool_calls_made: List[str] = []
        assistant_content = ""

        llm_client = LLMClient(self.oauth_client)

        try:
            async with llm_client:
                # Set context for intelligent file filtering
                self._agent.set_context(llm_client, user_message)

                # Track when assistant processing starts
                assistant_start_time = time.time()

                # Initial LLM call
                response = await llm_client.chat_completion(
                    messages=self._conversation_history,
                    tools=AGENT_TOOLS,
                )

                # Handle tool calls iteratively
                iteration = 0

                while iteration < self.MAX_ITERATIONS:
                    choice = response["choices"][0]
                    message = choice["message"]
                    finish_reason = choice["finish_reason"]

                    if finish_reason == "tool_calls" and "tool_calls" in message:
                        # Add assistant message with tool calls to history
                        tool_call_time = time.time()
                        assistant_msg = {
                            "role": "assistant",
                            "content": message.get("content") or "",
                            "tool_calls": cast(List[Dict[str, Any]], message["tool_calls"]),
                            "start_time": assistant_start_time,
                            "end_time": tool_call_time,
                        }
                        self._conversation_history.append(assistant_msg)
                        yield {"event": "message", "data": assistant_msg}

                        # Execute each tool call
                        for tool_call in message["tool_calls"]:
                            tool_name = tool_call["function"]["name"]
                            tool_calls_made.append(tool_name)
                            tool_exec_start = time.time()
                            result = await self._agent.execute_tool(tool_call)
                            tool_exec_end = time.time()

                            # Add tool result to history
                            tool_msg = {
                                "role": "tool",
                                "tool_call_id": tool_call["id"],
                                "content": result,
                                "start_time": tool_exec_start,
                                "end_time": tool_exec_end,
                            }
                            self._conversation_history.append(tool_msg)
                            yield {"event": "message", "data": tool_msg}

                        # Save state after tool calls
                        self._save_state()

                        # Reset start time for next assistant response
                        assistant_start_time = time.time()

                        # Get next response
                        response = await llm_client.chat_completion(
                            messages=self._conversation_history,
                            tools=AGENT_TOOLS,
                        )
                        iteration += 1
                    else:
                        # No more tool calls from LLM - exit the loop
                        break

                # Process final response
                choice = response["choices"][0]
                message = choice["message"]
                assistant_content = message.get("content") or ""

                # Handle empty responses
                if not assistant_content.strip():
                    for hist_msg in reversed(self._conversation_history):
                        if hist_msg.get("role") == "assistant" and hist_msg.get("content", "").strip():
                            assistant_content = hist_msg["content"]
                            break

                    if not assistant_content.strip() and tool_calls_made:
                        assistant_content = f"Done. Used tools: {', '.join(tool_calls_made)}"

                    if not assistant_content.strip():
                        yield {"event": "error", "data": {"error": "The assistant returned an empty response."}}
                        self._save_state()
                        return

                # Check if response seems incomplete (intent stated but no action taken)
                if iteration < self.MAX_ITERATIONS - 1:
                    seems_incomplete = await self._check_response_incomplete(
                        llm_client, assistant_content
                    )

                    if seems_incomplete:
                        # Add assistant message and prompt to continue
                        incomplete_end_time = time.time()
                        incomplete_msg = {
                            "role": "assistant",
                            "content": assistant_content,
                            "start_time": assistant_start_time,
                            "end_time": incomplete_end_time,
                        }
                        self._conversation_history.append(incomplete_msg)
                        yield {"event": "message", "data": incomplete_msg}

                        # Add a follow-up to prompt continuation (marked as auto-continue)
                        auto_continue_time = time.time()
                        auto_continue_msg = {
                            "role": "user",
                            "content": "[AUTO_CONTINUE] Please proceed with the action you described.",
                            "start_time": auto_continue_time,
                            "end_time": auto_continue_time,
                        }
                        self._conversation_history.append(auto_continue_msg)

                        # Reset start time for continuation response
                        assistant_start_time = time.time()

                        # Get continuation response
                        response = await llm_client.chat_completion(
                            messages=self._conversation_history,
                            tools=AGENT_TOOLS,
                        )

                        # Process continuation - may have tool calls
                        while iteration < self.MAX_ITERATIONS:
                            choice = response["choices"][0]
                            message = choice["message"]
                            finish_reason = choice["finish_reason"]

                            if finish_reason == "tool_calls" and "tool_calls" in message:
                                tool_call_time = time.time()
                                assistant_msg = {
                                    "role": "assistant",
                                    "content": message.get("content") or "",
                                    "tool_calls": cast(List[Dict[str, Any]], message["tool_calls"]),
                                    "start_time": assistant_start_time,
                                    "end_time": tool_call_time,
                                }
                                self._conversation_history.append(assistant_msg)
                                yield {"event": "message", "data": assistant_msg}

                                for tool_call in message["tool_calls"]:
                                    tool_name = tool_call["function"]["name"]
                                    tool_calls_made.append(tool_name)
                                    tool_exec_start = time.time()
                                    result = await self._agent.execute_tool(tool_call)
                                    tool_exec_end = time.time()

                                    tool_msg = {
                                        "role": "tool",
                                        "tool_call_id": tool_call["id"],
                                        "content": result,
                                        "start_time": tool_exec_start,
                                        "end_time": tool_exec_end,
                                    }
                                    self._conversation_history.append(tool_msg)
                                    yield {"event": "message", "data": tool_msg}

                                self._save_state()
                                assistant_start_time = time.time()
                                response = await llm_client.chat_completion(
                                    messages=self._conversation_history,
                                    tools=AGENT_TOOLS,
                                )
                                iteration += 1
                            else:
                                break

                        # Get the final content from continuation
                        assistant_content = response["choices"][0]["message"].get("content") or ""

                # Save final response
                final_end_time = time.time()
                final_msg = {
                    "role": "assistant",
                    "content": assistant_content,
                    "start_time": assistant_start_time,
                    "end_time": final_end_time,
                }
                self._conversation_history.append(final_msg)
                yield {"event": "message", "data": final_msg}

            # Save state after processing
            self._save_state()

            # Track successful chat message event
            user_id = await self.oauth_client.get_user_id()
            amplitude_service.track_chat_event(
                event_type=ChatEvent.CHAT_MESSAGE_SUBMITTED,
                user_id=user_id,
                project_path=self._project_path,
                start_time=user_timestamp,
                tool_calls_made=tool_calls_made,
                success=True,
            )

            # Yield done event with metadata
            yield {
                "event": "done",
                "data": {
                    "staged_changes": self._agent.get_staged_changes(),
                    "tool_calls_made": tool_calls_made,
                }
            }
        except Exception as e:
            from .exceptions import APIError

            # Check if this is an auth error (token expired/invalid)
            is_auth_error = False
            error_msg = str(e)
            if isinstance(e, APIError) and e.status_code in (401, 403):
                is_auth_error = True
                error_msg = "Your session has expired. Please sign in again."

            # Track failed chat message event
            user_id = await self.oauth_client.get_user_id()
            amplitude_service.track_chat_event(
                event_type=ChatEvent.CHAT_MESSAGE_SUBMITTED,
                user_id=user_id,
                project_path=self._project_path,
                start_time=user_timestamp,
                tool_calls_made=tool_calls_made,
                success=False,
                error_message=error_msg,
            )

            # Remove the user message we just added since the request failed
            if self._conversation_history and self._conversation_history[-1].get("role") == "user":
                self._conversation_history.pop()

            self._save_state()
            yield {"event": "error", "data": {"error": error_msg, "auth_error": is_auth_error}}

    async def get_staged_changes(self, data: Any) -> ChatRPCResponse:
        """Get current staged changes.

        Args:
            data: Not used

        Returns:
            Response with staged changes
        """
        if not self._agent:
            return {"result": {"edits": []}}
        return {"result": self._agent.get_staged_changes()}

    async def apply_changes(self, data: Any) -> ChatRPCResponse:
        """Apply all staged changes.

        Args:
            data: Not used

        Returns:
            Response with list of applied files and their status
        """
        if not self._agent:
            return {"error": "Chat not initialized"}

        results = self._agent.apply_staged_changes()

        # Build status message
        success_count = sum(1 for _, success, _ in results if success)
        fail_count = sum(1 for _, success, _ in results if not success)
        status_msg = f"Applied {success_count} file change(s)."
        if fail_count > 0:
            status_msg += f" Failed to apply {fail_count} change(s)."

        # Add to conversation history so it persists
        apply_time = time.time()
        self._conversation_history.append({
            "role": "assistant",
            "content": status_msg,
            "start_time": apply_time,
            "end_time": apply_time,
        })

        # Save state after applying (staged changes are now empty)
        self._save_state()

        return {
            "result": {
                "applied": [{"file_path": fp, "success": success, "message": msg} for fp, success, msg in results]
            }
        }

    async def clear_changes(self, data: Any) -> ChatRPCResponse:
        """Clear all staged changes.

        Args:
            data: Not used

        Returns:
            Response confirming changes were cleared
        """
        if self._agent:
            self._agent.clear_staged_changes()
            self._save_state()
        return {"result": {"cleared": True}}

    async def revert_change(self, data: Any) -> ChatRPCResponse:
        """Revert a specific staged change.

        Args:
            data: Dict with "file_path" key

        Returns:
            Response indicating whether revert was successful
        """
        if not self._agent:
            return {"error": "Chat not initialized"}

        file_path = data.get("file_path") if data else None
        if not file_path:
            return {"error": "file_path is required"}

        success = self._agent.revert_staged_change(file_path)
        if success:
            self._save_state()
        return {"result": {"reverted": success}}

    async def add_frontend_log(self, data: Any) -> ChatRPCResponse:
        """Add a frontend log entry (called from browser).

        Args:
            data: Dict with "log" key containing log line

        Returns:
            Response confirming log was added
        """
        if self._agent and data:
            log_line = data.get("log", "")
            if log_line:
                self._agent.add_frontend_log(log_line)
        return {"result": {"added": True}}

    async def authenticate(self, data: Any) -> ChatRPCResponse:
        """Start device authentication flow.

        Args:
            data: Not used

        Returns:
            Response with device authorization details
        """
        device_auth = await self.oauth_client.request_device_authorization()
        return {"result": device_auth}

    async def auth_poll(self, data: Any) -> ChatRPCResponse:
        """Poll for authentication completion.

        Args:
            data: Dict with "device_code" key

        Returns:
            Response indicating auth status
        """
        if not data:
            return {"error": "device_code is required"}

        device_code = data.get("device_code")
        if not device_code:
            return {"error": "device_code is required"}

        status_code, response = await self.oauth_client.check_authentication_status(device_code)

        if status_code == 200:
            await self.oauth_client._save_credentials(dict(response))
            # Reinitialize after successful auth
            return await self.initialize(None)
        else:
            error = response.get("error", "unknown_error")
            if error == "authorization_pending":
                return {"result": {"pending": True}}
            elif error == "slow_down":
                delay = 1 + (data.get("delayed", 0) if data else 0)
                return {"result": {"pending": True, "delay": delay}}
            elif error == "expired_token":
                return {"result": {"try_again": True}}
            else:
                return {"error": "Authentication failed"}

    async def logout(self, data: Any) -> ChatRPCResponse:
        """Log out and clear credentials.

        Args:
            data: Not used

        Returns:
            Response confirming logout
        """
        await self.oauth_client.logout()
        self._agent = None
        self._conversation_history = []
        # Clear the state file on logout
        self._clear_state()
        return {"result": {"logged_out": True}}

    async def reset_conversation(self, data: Any) -> ChatRPCResponse:
        """Reset the conversation history while keeping authentication.

        Args:
            data: Not used

        Returns:
            Response confirming reset
        """
        # Keep the agent but reset conversation
        self._conversation_history = [
            {
                "role": "system",
                "content": self._get_system_prompt(),
                "cache_control": {"type": "ephemeral"},
            }
        ]
        if self._agent:
            self._agent.clear_staged_changes()
        # Clear the state file (start fresh)
        self._clear_state()
        return {"result": {"reset": True}}
