"""V15: Auto-continue - fix Button import detection + return matching."""

import os

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

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

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

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

        default:"""
    c = c.replace("\n        default:", "\n" + block)
    with open(SSE, 'w') as f: f.write(c)
ok("sse.ts")

# 3. useAgentChat.ts
with open(HOOK) as f:
    lines = f.readlines()

new_lines = []
state_ok = hdl_ok = efx_ok = ret_ok = False

for line in lines:
    # Import useState
    if "import { useCallback, useEffect, useMemo, useRef } from 'react';" in line:
        line = line.replace("useRef }", "useRef, useState }")
    
    # State after callbacksRef
    if "callbacksRef.current = { onReady, onError, onSessionDead };" in line and not state_ok:
        new_lines.append(line)
        new_lines.append("  const [_showAc, _setShowAc] = useState(false);\n")
        new_lines.append("  const _acRef = useRef<Array<{ id: string; content: string; status: string }>>([]);\n")
        state_ok = True
        continue
    
    # Handler
    if "onSessionUpdate: (data) => {" in line and not hdl_ok:
        new_lines.append(line)
        new_lines.append("        if ((data as any).ac_plan) {\n")
        new_lines.append("          _acRef.current = (data as any).ac_plan;\n")
        new_lines.append("          _setShowAc(true);\n")
        new_lines.append("          return;\n")
        new_lines.append("        }\n")
        hdl_ok = True
        continue
    
    # Cancel + effect BEFORE 'return {'
    stripped = line.strip()
    if stripped == "return {" and not efx_ok:
        new_lines.append("  const _acCancel = useCallback(() => { _setShowAc(false); }, []);\n")
        new_lines.append("  useEffect(() => {\n")
        new_lines.append("    if (!_showAc) return;\n")
        new_lines.append("    const plan = _acRef.current;\n")
        new_lines.append("    const timer = setTimeout(() => {\n")
        new_lines.append("      const s = chatActionsRef.current.setMessages;\n")
        new_lines.append("      const m = chatActionsRef.current.messages;\n")
        new_lines.append("      if (s && plan.length > 0) {\n")
        new_lines.append("        const t = plan.map((i: { content: string }) => '- ' + i.content).join('\\\\n');\n")
        new_lines.append("        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")
        new_lines.append("        s([...m, nu]);\n")
        new_lines.append("      }\n")
        new_lines.append("      _setShowAc(false);\n")
        new_lines.append("    }, 10000);\n")
        new_lines.append("    return () => clearTimeout(timer);\n")
        new_lines.append("  }, [_showAc]);\n\n")
        efx_ok = True
    
    # Add return values after "refreshMessages,"
    if "refreshMessages," in line and not ret_ok:
        new_lines.append(line)
        new_lines.append("    _showAc,\n")
        new_lines.append("    _acCancel,\n")
        ret_ok = True
        continue
    
    new_lines.append(line)

with open(HOOK, 'w') as f:
    f.writelines(new_lines)
ok(f"useAgentChat.ts (s={state_ok} h={hdl_ok} e={efx_ok} r={ret_ok})")

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

# Check for exact Button in MUI import (not "IconButton" which contains "Button")
if "  Button,\n" not in content and "  Button\n" not in content:
    content = content.replace(
        "  Tooltip,",
        "  Tooltip,\n  Button,"
    )

# Add props to interface
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) {")

    # Add pause button before <JobsUpgradeDialog
    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")

with open(INPUT, 'w') as f:
    f.write(content)
ok("ChatInput.tsx")

# 5. 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} ")
        with open(SESS, 'w') as f: f.write(c)
    ok("SessionChat.tsx")

# 6. agent_loop.py
if os.path.exists(AGENT):
    with open(AGENT) as f: c = f.read()
    if "_unfinished_plan" not in c:
        fn = "\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"
        c = c.replace("class Handlers:", fn + "\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"
        )
        marker = "  # -- End of turn --"
        if marker in c:
            c = c.replace(marker, check + marker)
    with open(AGENT, 'w') as f: f.write(c)
    ok("agent_loop.py")
else:
    print("SKIP: agent_loop.py")

print("\nDONE: V15")