"""V16: Auto-continue - match exact 'return {\n    messages: chat.messages,' (only one in file)"""

import os, sys

F = lambda n: "/source/frontend/src/" + n
EVENTS = F("types/events.ts")
SSE = F("lib/sse-chat-transport.ts")
HOOK = F("hooks/useAgentChat.ts")
INPUT = F("components/Chat/ChatInput.tsx")
SESS = F("components/SessionChat.tsx")
AGENT = "/app/agent/core/agent_loop.py"

def ok(f): print(f"OK: {f}")

# --- events.ts ---
with open(EVENTS) as f: c = f.read()
if "'task_incomplete'" not in c:
    c = c.replace("  | 'interrupted'", "  | 'interrupted'\n  | 'task_incomplete'")
    open(EVENTS, 'w').write(c)
ok("events.ts")

# --- sse.ts ---
with open(SSE) as f: c = f.read()
if "case 'task_incomplete'" not in c:
    c = c.replace("\n        default:", "\n      case 'task_incomplete':\n          (sideChannel as any).onTaskIncomplete(\n            (event.data?.incomplete_plan as Array<{ id: string; content: string; status: string }>) || [],\n          );\n          break;\n\n        default:")
    open(SSE, 'w').write(c)
ok("sse.ts")

# --- useAgentChat.ts ---
with open(HOOK) as f: content = f.read()

# useState import
if "useState" not in content:
    content = content.replace("useRef }", "useRef, useState }")

# State after callbacksRef (insert once)
if "_showAc" not in content:
    s = ("  // Auto-continue state\n"
         "  const [_showAc, _setShowAc] = useState(false);\n"
         "  const _acRef = useRef<Array<{ id: string; content: string; status: string }>>([]);\n")
    content = content.replace(
        "callbacksRef.current = { onReady, onError, onSessionDead };",
        "callbacksRef.current = { onReady, onError, onSessionDead };\n" + s
    )

# Handler in onSessionUpdate
if "ac_plan" not in content:
    h = ("        // Auto-continue check\n"
         "        if ((data as any).ac_plan) {\n"
         "          _acRef.current = (data as any).ac_plan;\n"
         "          _setShowAc(true);\n"
         "          return;\n"
         "        }\n")
    content = content.replace("onSessionUpdate: (data) => {", "onSessionUpdate: (data) => {\n" + h)

# CRITICAL: Only match the exact two-line return that has messages: chat.messages
# Only ONE place in the file has this exact pattern
marker = "  return {\n    messages: chat.messages,"
if marker in content and "_acCancel" not in content:
    cancel_effect = (
        "  const _acCancel = useCallback(() => { _setShowAc(false); }, []);\n"
        "  useEffect(() => {\n"
        "    if (!_showAc) return;\n"
        "    const plan = _acRef.current;\n"
        "    const timer = setTimeout(() => {\n"
        "      const s = chatActionsRef.current.setMessages;\n"
        "      const m = chatActionsRef.current.messages;\n"
        "      if (s && plan.length > 0) {\n"
        "        const t = plan.map((i: { content: string }) => '- ' + i.content).join('\\\\n');\n"
        "        const nu = { id: 'ac-'+Date.now(), role: 'user' as const, parts: [{ type: 'text' as const, text: '[TIEP TUC] Task chua hoan thanh:\\\\n' + t }], content: '' };\n"
        "        s([...m, nu]);\n"
        "      }\n"
        "      _setShowAc(false);\n"
        "    }, 10000);\n"
        "    return () => clearTimeout(timer);\n"
        "  }, [_showAc]);\n\n"
    )
    content = content.replace(marker, cancel_effect + marker)
    
    # Add to return values
    content = content.replace(
        "refreshMessages,\n  };",
        "refreshMessages,\n    _showAc,\n    _acCancel,\n  };"
    )
    
with open(HOOK, 'w').write(content)
ok("useAgentChat.ts")

# --- ChatInput.tsx ---
with open(INPUT) as f: content = f.read()

# Button: check for standalone ",  Button" or "  Button" NOT "IconButton"
if ",  Button,\n" not in content and ",  Button\n" not in content:
    content = content.replace(
        "  Tooltip,",
        "  Tooltip,\n  Button,"
    )

if "_showAc" not in content:
    content = content.replace("placeholder?: string;", "placeholder?: string;\n  _showAc?: boolean;\n  _acCancel?: () => void;")
    content = content.replace("placeholder = 'Ask anything...' }: ChatInputProps) {",
                             "placeholder = 'Ask anything...', _showAc = false, _acCancel }: ChatInputProps) {")
    btn = ("        {_showAc && _acCancel && (\n"
           "          <Box sx={{ display: 'flex', justifyContent: 'center', mt: 1 }}>\n"
           "            <Button variant=\"outlined\" size=\"small\" onClick={_acCancel}\n"
           "              sx={{ textTransform: 'none', fontSize: '0.7rem', opacity: 0.7, '&:hover': { opacity: 1 } }}>\n"
           "              TAM_DUNG_TU_DONG_TIEP_TUC (10s)\n"
           "            </Button>\n"
           "          </Box>\n"
           "        )}\n")
    content = content.replace("<JobsUpgradeDialog", btn + "        <JobsUpgradeDialog")

open(INPUT, 'w').write((content).read() if hasattr(content, 'read') else content)
ok("ChatInput.tsx")

# --- SessionChat.tsx ---
if os.path.exists(SESS):
    with open(SESS) as f: c = f.read()
    if "_showAc" not in c:
        c = c.replace("<ChatInput ", "<ChatInput _showAc={_showAc} _acCancel={_acCancel} ")
        open(SESS, 'w').write(c)
    ok("SessionChat.tsx")

# --- agent_loop.py ---
if os.path.exists(AGENT):
    with open(AGENT) as f: c = f.read()
    if "_unfinished_plan" not in c:
        c = c.replace("class Handlers:", "\n\ndef _unfinished_plan(s):\n    p = getattr(s, 'current_plan', None) or []\n    return [it for it in p if it.get('status') in ('pending', 'in_progress')]\n\n\nclass Handlers:")
    if "ac_plan" not in c:
        check = ("\n            if not llm_result.tool_calls_acc and llm_result.content:\n"
                 "                unfinished = _unfinished_plan(session)\n"
                 "                if unfinished:\n"
                 "                    await session.send_event(Event(event_type='session_update', data={\n"
                 "                        'ac_plan': unfinished,\n"
                 "                    }))\n")
        c = c.replace("  # -- End of turn --", check + "  # -- End of turn --")
    open(AGENT, 'w').write(c)
    ok("agent_loop.py")
else:
    print("SKIP: agent_loop.py")

print("\nDONE: V16")