"""Patch script: Add auto-continue feature for free models in ml-intern.

Apply this script during Docker build to enable:
1. Detection when streaming ends but task incomplete
2. "Tự động tiếp tục (10s)" button - auto continues after 10s
3. "Tạm dừng" button - lets user send different content

Usage: Run this during Docker build after cloning source.
"""

import os
import re

# Paths - adjust if needed
USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"
SSE_TRANSPORT = "/source/frontend/src/lib/sse-chat-transport.ts"
AGENT_LOOP = "/app/agent/core/agent_loop.py"

# =============================================================================
# Patch 1: useAgentChat.ts
# =============================================================================
def patch_use_agent_chat():
    if not os.path.exists(USE_AGENT_CHAT):
        print(f"SKIP: {USE_AGENT_CHAT} not found")
        return
        
    with open(USE_AGENT_CHAT, "r", encoding="utf-8") as f:
        content = f.read()
    
    # Add useState to imports if missing
    if "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';" not in content:
        content = content.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # Add auto-continue state after callbacksRef
    state_code = '''
  // Auto-continue state for free models when task incomplete
  const autoContinueTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
  const [showAutoContinue, setShowAutoContinue] = useState(false);
  const [taskIncompleteInfo, setTaskIncompleteInfo] = useState<{
    incompletePlan: Array<{ id: string; content: string; status: string }>;
  } | null>(null);
'''
    if "autoContinueTimerRef" not in content and "callbacksRef.current = { onReady, onError, onSessionDead }" in content:
        content = content.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_code
        )
    
    # Add onTaskIncomplete to SideChannelCallbacks interface
    if "onTaskIncomplete" not in content:
        content = content.replace(
            "onInterrupted: () => void;",
            "onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    with open(USE_AGENT_CHAT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {USE_AGENT_CHAT}")

# =============================================================================
# Patch 2: sse-chat-transport.ts  
# =============================================================================
def patch_sse_transport():
    if not os.path.exists(SSE_TRANSPORT):
        print(f"SKIP: {SSE_TRANSPORT} not found")
        return
    
    with open(SSE_TRANSPORT, "r", encoding="utf-8") as f:
        content = f.read()
    
    # Add task_incomplete case
    if "case 'task_incomplete'" not in content:
        task_incomplete_case = '''      case 'task_incomplete':
          sideChannel.onTaskIncomplete(
            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],
          );
          break;
'''
        # Insert before turn_complete case
        content = content.replace(
            "case 'turn_complete':",
            task_incomplete_case + "\n        case 'turn_complete':"
        )
    
    with open(SSE_TRANSPORT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {SSE_TRANSPORT}")

# =============================================================================
# Patch 3: agent_loop.py
# =============================================================================
def patch_agent_loop():
    if not os.path.exists(AGENT_LOOP):
        print(f"SKIP: {AGENT_LOOP} not found")
        return
    
    with open(AGENT_LOOP, "r", encoding="utf-8") as f:
        content = f.read()
    
    # Add check function if not exists
    check_func = '''
def _check_task_incomplete(session: Session, llm_result: LLMResult) -> bool:
    """Check if model stopped but task is incomplete (has unfinished plan items)."""
    if llm_result.tool_calls_acc:
        return False  # Has tool calls, not incomplete
    plan = getattr(session, "current_plan", None) or []
    unfinished = [item for item in plan if item.get("status") in ("pending", "in_progress")]
    return len(unfinished) > 0
'''
    if "_check_task_incomplete" not in content:
        # Insert after LLMResult dataclass
        content = content.replace(
            "class Handlers:",
            check_func + "\n\nclass Handlers:"
        )
    
    with open(AGENT_LOOP, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Patched {AGENT_LOOP}")


if __name__ == "__main__":
    patch_use_agent_chat()
    patch_sse_transport()
    patch_agent_loop()
    print("\nDONE: Auto-continue patches applied")