"""Patch useAgentChat.ts - auto-continue feature."""

USE_AGENT_CHAT = "/source/frontend/src/hooks/useAgentChat.ts"

def patch():
    import os
    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()
    
    # 1. Add useState to imports
    if "useState" not in content:
        content = content.replace(
            "import { useCallback, useEffect, useMemo, useRef } from 'react';",
            "import { useCallback, useEffect, useMemo, useRef, useState } from 'react';"
        )
    
    # 2. Add state declarations after callbacksRef.current = { ... }
    state_decl = '''
  // 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:
        content = content.replace(
            "callbacksRef.current = { onReady, onError, onSessionDead };",
            "callbacksRef.current = { onReady, onError, onSessionDead };" + state_decl
        )
    
    # 3. Add onTaskIncomplete to SideChannelCallbacks interface
    if "onTaskIncomplete:" not in content and "onInterrupted:" in content:
        content = content.replace(
            "  onInterrupted: () => void;",
            "  onInterrupted: () => void;\n  onTaskIncomplete: (incompletePlan: Array<{ id: string; content: string; status: string }>) => void;"
        )
    
    # 4. Add onTaskIncomplete handler in sideChannel
    sidechannel_handler = '''      onTaskIncomplete: (incompletePlan) => {
        setTaskIncompleteInfo({ incompletePlan });
        setShowAutoContinue(true);
        // Auto-start 10s timer for free models
        if (autoContinueTimerRef.current) {
          clearTimeout(autoContinueTimerRef.current);
        }
        autoContinueTimerRef.current = setTimeout(() => {
          startAutoContinue();
        }, 10000);
      },'''
    
    if "onTaskIncomplete: (incompletePlan)" not in content:
        content = content.replace(
            "      onInterrupted: () => { /* no-op — handled by stop() caller */ },",
            sidechannel_handler + "\n      onInterrupted: () => { /* no-op — handled by stop() caller */ },"
        )
    
    # 5. Add auto-continue functions before return
    auto_funcs = '''
  // -- Auto-continue for free models when task incomplete -----------------
  const startAutoContinue = useCallback(() => {
    setShowAutoContinue(false);
    setTaskIncompleteInfo(null);
    if (autoContinueTimerRef.current) {
      clearTimeout(autoContinueTimerRef.current);
      autoContinueTimerRef.current = null;
    }
    // Build continuation message from incomplete plan
    const incompleteItems = taskIncompleteInfo?.incompletePlan || [];
    const continuationText = incompleteItems.length > 0
      ? `Tiếp tục từ các task chưa hoàn thành:\\n${incompleteItems.map(i => `- ${i.content}`).join('\\n')}`
      : 'Tiếp tục nhiệm vụ.';
    chat.sendMessage({
      text: `[TỰ ĐỘNG TIẾP TỤC] ${continuationText}`,
      metadata: { createdAt: new Date().toISOString() },
    });
  }, [taskIncompleteInfo, chat]);

  const cancelAutoContinue = useCallback(() => {
    setShowAutoContinue(false);
    setTaskIncompleteInfo(null);
    if (autoContinueTimerRef.current) {
      clearTimeout(autoContinueTimerRef.current);
      autoContinueTimerRef.current = null;
    }
  }, []);

  // Cleanup timer on unmount
  useEffect(() => {
    return () => {
      if (autoContinueTimerRef.current) {
        clearTimeout(autoContinueTimerRef.current);
      }
    };
  }, []);
'''
    
    if "startAutoContinue" not in content:
        content = content.replace(
            '\n  return {\n    messages: chat.messages,',
            auto_funcs + '\n  return {\n    messages: chat.messages,'
        )
    
    # 6. Update return to include new values
    if "showAutoContinue," not in content:
        content = content.replace(
            "refreshMessages,\n  };",
            "refreshMessages,\n    showAutoContinue,\n    taskIncompleteInfo,\n    startAutoContinue,\n    cancelAutoContinue,\n  };"
        )
    
    with open(USE_AGENT_CHAT, "w", encoding="utf-8") as f:
        f.write(content)
    print(f"OK: Complete patch applied to {USE_AGENT_CHAT}")

if __name__ == "__main__":
    patch()