"""Shared browser UI components for KISS agent viewers."""

import json
import logging
import queue
import socket
import threading
import time
from typing import Any

import yaml

from kiss.core.printer import Printer, extract_extras, extract_path_and_lang, truncate_result

logger = logging.getLogger(__name__)

def find_free_port() -> int:
    """Find and return an available TCP port on localhost.

    Returns:
        int: A free port number.
    """
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return int(s.getsockname()[1])


BASE_CSS = r"""
*{margin:0;padding:0;box-sizing:border-box}
html{height:100%;overflow:hidden}
:root{
  --bg:#0d1117;--surface:#161b22;--surface2:#1c2128;--border:#30363d;
  --text:#e6edf3;--dim:#8b949e;--accent:#58a6ff;--green:#3fb950;
  --red:#f85149;--yellow:#d29922;--cyan:#79c0ff;--purple:#bc8cff;
}
body{
  font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Helvetica,Arial,sans-serif;
  background:var(--bg);color:var(--text);line-height:1.6;
  height:100%;display:flex;flex-direction:column;overflow:hidden;
}
header{
  background:linear-gradient(135deg,var(--surface) 0%,var(--surface2) 100%);
  border-bottom:1px solid var(--border);padding:12px 24px;
  display:flex;align-items:center;justify-content:space-between;flex-shrink:0;
}
.logo{font-size:18px;font-weight:700;color:var(--accent);letter-spacing:-.3px}
.logo span{color:var(--dim);font-weight:400;font-size:14px;margin-left:8px}
.status{display:flex;align-items:center;gap:8px;font-size:13px;color:var(--dim)}
.dot{width:8px;height:8px;border-radius:50%;background:var(--dim)}
.dot.running{background:var(--green);animation:pulse 2s infinite}
.dot.done{animation:none;background:var(--dim)}
@keyframes pulse{0%,100%{opacity:1}50%{opacity:.3}}
"""

OUTPUT_CSS = r"""
.ev{margin-bottom:6px;animation:fadeIn .15s ease}
@keyframes fadeIn{from{opacity:0;transform:translateY(3px)}to{opacity:1;transform:none}}
.think{
  padding:10px 16px;margin:10px 0;
  border:1px solid var(--border);border-radius:8px;
  background:rgba(121,192,255,.04);max-height:200px;overflow-y:auto;
}
.think .lbl{
  font-size:11px;font-weight:600;text-transform:uppercase;
  letter-spacing:.06em;color:var(--cyan);margin-bottom:4px;
  display:flex;align-items:center;gap:6px;cursor:pointer;user-select:none;
}
.think .lbl .arrow{transition:transform .2s;display:inline-block}
.think .lbl .arrow.collapsed{transform:rotate(-90deg)}
.think .cnt{
  font-size:11px;color:var(--dim);font-style:italic;
  white-space:pre-wrap;word-break:break-word;
}
.think .cnt.hidden{display:none}
.txt{font-size:8px;white-space:pre-wrap;word-break:break-word;padding:2px 0;line-height:1.4}
.tc{
  border:1px solid var(--border);border-radius:8px;margin:10px 0;
  overflow:hidden;background:var(--surface);transition:box-shadow .2s;
}
.tc:hover{box-shadow:0 2px 12px rgba(0,0,0,.3)}
.tc-h{
  padding:9px 14px;background:var(--surface2);
  display:flex;align-items:center;gap:10px;cursor:pointer;user-select:none;
}
.tc-h:hover{background:rgba(48,54,61,.8)}
.tc-h .chv{color:var(--dim);transition:transform .2s;font-size:11px;flex-shrink:0}
.tc-h .chv.open{transform:rotate(90deg)}
.tn{font-weight:600;font-size:13px;color:var(--accent)}
.tp{font-size:12px;color:var(--cyan);font-family:'SF Mono','Fira Code',monospace}
.td{font-size:11px;color:var(--dim);font-style:italic}
.tc-b{
  padding:10px 14px;max-height:300px;overflow-y:auto;
  font-family:'SF Mono','Fira Code',monospace;font-size:11px;line-height:1.5;
}
.tc-b.hide{display:none}
.tc-b pre{margin:4px 0;white-space:pre-wrap;word-break:break-word}
.diff-old{
  color:var(--red);background:rgba(248,81,73,.08);
  padding:1px 6px;display:block;white-space:pre-wrap;word-break:break-word;
}
.diff-new{
  color:var(--green);background:rgba(63,185,80,.08);
  padding:1px 6px;display:block;white-space:pre-wrap;word-break:break-word;
}
.diff-ctx{
  color:var(--dim);padding:1px 6px;display:block;opacity:0.5;
  white-space:pre-wrap;word-break:break-word;
}
.diff-hl-del{background:rgba(248,81,73,.3);border-radius:2px;padding:0 1px}
.diff-hl-add{background:rgba(63,185,80,.3);border-radius:2px;padding:0 1px}
.extra{color:var(--dim);margin:2px 0}
.tr{
  padding:8px 14px;margin:6px 0;
  border:1px solid var(--border);border-radius:8px;
  font-family:'SF Mono','Fira Code',monospace;
  font-size:11px;max-height:200px;overflow-y:auto;
  white-space:pre-wrap;word-break:break-word;background:rgba(63,185,80,.04);
}
.tr.err{background:rgba(248,81,73,.04);border-color:rgba(248,81,73,.2)}
.tr .rl{
  font-size:11px;font-weight:600;text-transform:uppercase;
  letter-spacing:.06em;margin-bottom:4px;
}
.tr .rl.ok{color:var(--green)}
.tr .rl.fail{color:var(--red)}
.rc{
  border:1px solid var(--green);border-radius:8px;
  margin:20px 0;overflow:hidden;background:var(--surface);
}
.rc-h{
  padding:14px 20px;background:rgba(63,185,80,.08);
  display:flex;align-items:center;justify-content:space-between;
}
.rc-h h3{color:var(--green);font-size:14px;font-weight:600}
.rs{font-size:12px;color:var(--dim);display:flex;gap:18px}
.rs b{color:var(--text);font-weight:500}
.rc-body{
  padding:16px 20px;font-size:13px;max-height:400px;overflow-y:auto;
  word-break:break-word;line-height:1.7;
}
.rc-body.pre{white-space:pre-wrap}
.prompt{
  border:1px solid var(--border);border-radius:8px;margin:10px 0;
  overflow:hidden;background:var(--surface);
}
.prompt-h{
  padding:8px 16px;background:rgba(121,192,255,.08);
  font-size:11px;font-weight:600;color:var(--cyan);
  text-transform:uppercase;letter-spacing:.04em;
}
.prompt-body{
  padding:12px 16px;font-size:13px;white-space:pre-wrap;
  word-break:break-word;line-height:1.6;max-height:400px;overflow-y:auto;
}
.sys{
  font-size:13px;color:var(--dim);font-family:'SF Mono','Fira Code',monospace;
  white-space:pre-wrap;word-break:break-word;padding:2px 0;
}
.bash-panel{
  background:#08080a;border:1px solid var(--border);
  border-radius:8px;margin:2px 0 8px;padding:10px 12px;
  max-height:300px;overflow-y:auto;
  font-family:'SF Mono','Fira Code','Cascadia Code',monospace;
  font-size:8px;line-height:1.5;color:rgba(255,255,255,0.65);
  white-space:pre-wrap;word-break:break-word;
}
.usage{
  border:1px solid var(--border);border-radius:8px;margin:6px 0;
  padding:4px 12px;background:var(--surface);font-size:11px;
  color:var(--dim);font-style:italic;
  font-family:'SF Mono','Fira Code',monospace;
  white-space:nowrap;overflow-x:auto;
}
.empty-msg{
  text-align:center;color:var(--dim);
  padding:80px 20px;font-size:15px;line-height:2;
}
.spinner{
  display:flex;align-items:center;gap:10px;
  padding:16px 0;color:var(--dim);font-size:13px;
  animation:fadeIn .3s ease;
}
.spinner::before{
  content:'';width:16px;height:16px;
  border:2px solid var(--border);border-top-color:var(--accent);
  border-radius:50%;flex-shrink:0;animation:spin .8s linear infinite;
}
@keyframes spin{to{transform:rotate(360deg)}}
::-webkit-scrollbar{width:6px}
::-webkit-scrollbar-track{background:transparent}
::-webkit-scrollbar-thumb{background:var(--border);border-radius:3px}
::-webkit-scrollbar-thumb:hover{background:var(--dim)}
code{font-family:'SF Mono','Fira Code',monospace}
.hljs{background:var(--surface2)!important;border-radius:6px;padding:10px!important}
"""

EVENT_HANDLER_JS = r"""
function esc(t){var d=document.createElement('div');d.textContent=t;return d.innerHTML}
function mkEl(tag,cls){var e=document.createElement(tag);if(cls)e.className=cls;return e}
function toggleTC(el){
  el.nextElementSibling.classList.toggle('hide');
  el.querySelector('.chv').classList.toggle('open');
}
function toggleThink(el){
  var p=el.parentElement;
  p.querySelector('.cnt').classList.toggle('hidden');
  el.querySelector('.arrow').classList.toggle('collapsed');
}
function lineDiff(a,b){
  var al=a.split('\n'),bl=b.split('\n'),m=al.length,n=bl.length;
  var dp=[];
  for(var i=0;i<=m;i++){dp[i]=new Array(n+1);dp[i][0]=0;}
  for(var j=0;j<=n;j++)dp[0][j]=0;
  for(var i=1;i<=m;i++)for(var j=1;j<=n;j++)
    dp[i][j]=al[i-1]===bl[j-1]?dp[i-1][j-1]+1:Math.max(dp[i-1][j],dp[i][j-1]);
  var ops=[],i=m,j=n;
  while(i>0||j>0){
    if(i>0&&j>0&&al[i-1]===bl[j-1]){ops.unshift({t:'=',o:al[--i],n:bl[--j]});}
    else if(j>0&&(i===0||dp[i][j-1]>=dp[i-1][j])){ops.unshift({t:'+',n:bl[--j]});}
    else{ops.unshift({t:'-',o:al[--i]});}
  }
  return ops;
}
function hlInline(oldL,newL){
  var mn=Math.min(oldL.length,newL.length),pre=0,suf=0;
  while(pre<mn&&oldL[pre]===newL[pre])pre++;
  while(suf<mn-pre&&oldL[oldL.length-1-suf]===newL[newL.length-1-suf])suf++;
  var pf=oldL.substring(0,pre),sf=suf?oldL.substring(oldL.length-suf):'';
  return{
    o:esc(pf)+'<span class="diff-hl-del">'+esc(oldL.substring(pre,oldL.length-suf))
      +'</span>'+esc(sf),
    n:esc(pf)+'<span class="diff-hl-add">'+esc(newL.substring(pre,newL.length-suf))
      +'</span>'+esc(sf)
  };
}
function renderDiff(oldStr,newStr){
  var ops=lineDiff(oldStr,newStr),html='',i=0;
  while(i<ops.length){
    var dels=[],adds=[];
    while(i<ops.length&&ops[i].t==='-'){dels.push(ops[i++]);}
    while(i<ops.length&&ops[i].t==='+'){adds.push(ops[i++]);}
    if(dels.length||adds.length){
      var pairs=Math.min(dels.length,adds.length);
      for(var p=0;p<pairs;p++){
        var h=hlInline(dels[p].o,adds[p].n);
        html+='<div class="diff-old">- '+h.o+'</div>';
        html+='<div class="diff-new">+ '+h.n+'</div>';
      }
      for(var p=pairs;p<dels.length;p++)
        html+='<div class="diff-old">- '+esc(dels[p].o)+'</div>';
      for(var p=pairs;p<adds.length;p++)
        html+='<div class="diff-new">+ '+esc(adds[p].n)+'</div>';
      continue;
    }
    html+='<div class="diff-ctx">  '+esc(ops[i].o)+'</div>';i++;
  }
  return html;
}
function handleOutputEvent(ev,O,state){
  var t=ev.type;
  switch(t){
  case'thinking_start':
    state.thinkEl=mkEl('div','ev think');
    state.thinkEl.innerHTML=
      '<div class="lbl" onclick="toggleThink(this)">'
      +'<span class="arrow">\u25BE</span> Thinking</div>'
      +'<div class="cnt"></div>';
    O.appendChild(state.thinkEl);break;
  case'thinking_delta':
    if(state.thinkEl){
      var tc=state.thinkEl.querySelector('.cnt');
      tc.textContent+=(ev.text||'').replace(/\n\n+/g,'\n');
      state.thinkEl.scrollTop=state.thinkEl.scrollHeight;
    }break;
  case'thinking_end':
    if(state.thinkEl){
      state.thinkEl.querySelector('.lbl').innerHTML=
        '<span class="arrow collapsed">\u25BE</span> Thinking (click to expand)';
      state.thinkEl.querySelector('.cnt').classList.add('hidden');
    }
    state.thinkEl=null;break;
  case'text_delta':
    if(!state.txtEl){state.txtEl=mkEl('div','txt');O.appendChild(state.txtEl)}
    state.txtEl.textContent+=(ev.text||'').replace(/\n\n+/g,'\n');break;
  case'text_end':state.txtEl=null;break;
  case'tool_call':{
    if(state.bashPanel&&state.bashBuf){state.bashPanel.textContent+=state.bashBuf;state.bashBuf=''}
    state.bashPanel=null;state.bashRaf=0;
    var c=mkEl('div','ev tc');
    var h='<span class="chv open">\u25B6</span><span class="tn">'+esc(ev.name)+'</span>';
    if(ev.path){var ep=esc(ev.path).replace(/"/g,'&quot;');
      h+='<span class="tp" data-path="'+ep+'"> '+esc(ev.path)+'</span>';}
    if(ev.description)h+='<span class="td"> '+esc(ev.description)+'</span>';
    var b='';
    if(ev.command)b+='<pre><code class="language-bash">'+esc(ev.command)+'</code></pre>';
    if(ev.content){
      var lc=ev.lang?'language-'+esc(ev.lang):'';
      b+='<pre><code class="'+lc+'">'+esc(ev.content)+'</code></pre>';
    }
    if(ev.old_string!==undefined&&ev.new_string!==undefined){
      b+=renderDiff(ev.old_string,ev.new_string);
    }else{
      if(ev.old_string!==undefined)b+='<div class="diff-old">- '+esc(ev.old_string)+'</div>';
      if(ev.new_string!==undefined)b+='<div class="diff-new">+ '+esc(ev.new_string)+'</div>';
    }
    if(ev.extras){for(var k in ev.extras)
      b+='<div class="extra">'+esc(k)+': '+esc(ev.extras[k])+'</div>'}
    var body=b||'<em style="color:var(--dim)">No arguments</em>';
    c.innerHTML='<div class="tc-h" onclick="toggleTC(this)">'+h+'</div>'
      +'<div class="tc-b'+(b?'':' hide')+'">'+body+'</div>';
    O.appendChild(c);
    if(ev.command){var bp=mkEl('div','bash-panel');O.appendChild(bp);state.bashPanel=bp}
    if(typeof hljs!=='undefined')c.querySelectorAll('pre code').forEach(
      function(bl){hljs.highlightElement(bl)});
    break}
  case'tool_result':{
    if(state.bashPanel&&state.bashBuf){state.bashPanel.textContent+=state.bashBuf;state.bashBuf=''}
    var hadBash=!!state.bashPanel;
    state.bashPanel=null;state.bashRaf=0;
    if(hadBash&&!ev.is_error)break;
    var r=mkEl('div','ev tr'+(ev.is_error?' err':''));
    var lb=ev.is_error?'FAILED':'OK';
    var lc2=ev.is_error?'fail':'ok';
    r.innerHTML='<div class="rl '+lc2+'">'+lb+'</div>'+esc(ev.content);
    O.appendChild(r);break}
  case'system_output':{
    if(state.bashPanel){
      if(!state.bashBuf)state.bashBuf='';
      state.bashBuf+=(ev.text||'');
      if(!state.bashRaf){state.bashRaf=requestAnimationFrame(function(){
        if(state.bashPanel)state.bashPanel.textContent+=state.bashBuf;
        state.bashBuf='';state.bashRaf=0;
        if(state.bashPanel)state.bashPanel.scrollTop=state.bashPanel.scrollHeight;
      })}
    }else{
      var s=mkEl('div','ev sys');
      s.textContent=(ev.text||'').replace(/\n\n+/g,'\n');
      O.appendChild(s);
    }break}
  case'result':{
    var rc=mkEl('div','ev rc');
    var rb='';
    if(ev.success===false){
      var fs='color:var(--red);font-weight:700;font-size:16px;margin-bottom:12px';
      rb+='<div style="'+fs+'">Status: FAILED</div>';
    }
    var usePre=true;
    if(ev.summary){
      var sum=(ev.summary||'').replace(/\n{3,}/g,'\n\n').trim();
      if(typeof marked!=='undefined'){rb+=marked.parse(sum);usePre=false;}
      else{rb+=esc(sum);}
    }else{
      rb+=esc((ev.text||'(no result)').replace(/\n{3,}/g,'\n\n').trim());
    }
    rc.innerHTML='<div class="rc-h"><h3>Result</h3><div class="rs">'
      +'<span>Steps <b>'+(ev.step_count||0)+'</b></span>'
      +'<span>Tokens <b>'+(ev.total_tokens||0)+'</b></span>'
      +'<span>Cost <b>'+(ev.cost||'N/A')+'</b></span>'
      +'</div></div><div class="rc-body'
      +(usePre?' pre':'')+'">'+rb+'</div>';
    if(typeof hljs!=='undefined')rc.querySelectorAll('pre code').forEach(
      function(bl){hljs.highlightElement(bl)});
    O.appendChild(rc);break}
  case'prompt':{
    var p=mkEl('div','ev prompt');
    p.innerHTML='<div class="prompt-h">Prompt</div>'
      +'<div class="prompt-body">'+esc(ev.text||'')+'</div>';
    O.appendChild(p);break}
  case'usage_info':{
    var u=mkEl('div','ev usage');
    u.textContent=ev.text||'';
    O.appendChild(u);break}
  }
}
"""

HTML_HEAD = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{title}</title>
<link rel="stylesheet"
  href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github-dark.min.css">
<script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/marked/12.0.1/marked.min.js"></script>
<style>
{css}
</style>
</head>
"""


class BaseBrowserPrinter(Printer):
    def __init__(self) -> None:
        self._clients: list[queue.Queue[dict[str, Any]]] = []
        self._lock = threading.Lock()
        self._current_block_type = ""
        self._tool_name = ""
        self._tool_json_buffer = ""
        self._bash_lock = threading.Lock()
        self._bash_buffer: list[str] = []
        self._bash_last_flush = 0.0
        self._bash_flush_timer: threading.Timer | None = None
        self.stop_event = threading.Event()

    def reset(self) -> None:
        """Reset internal streaming and tool-parsing state for a new turn."""
        self._current_block_type = ""
        self._tool_name = ""
        self._tool_json_buffer = ""
        with self._bash_lock:
            self._bash_buffer.clear()
            if self._bash_flush_timer is not None:
                self._bash_flush_timer.cancel()
                self._bash_flush_timer = None

    def _flush_bash(self) -> None:
        with self._bash_lock:
            if self._bash_flush_timer is not None:
                self._bash_flush_timer.cancel()
                self._bash_flush_timer = None
            if self._bash_buffer:
                text = "".join(self._bash_buffer)
                self._bash_buffer.clear()
                self._bash_last_flush = time.monotonic()
            else:
                text = ""
        if text:
            self.broadcast({"type": "system_output", "text": text})

    @staticmethod
    def _parse_result_yaml(raw: str) -> dict[str, Any] | None:
        try:
            data = yaml.safe_load(raw)
        except Exception:
            logger.debug("Exception caught", exc_info=True)
            return None
        if isinstance(data, dict) and "summary" in data:
            return data
        return None

    def broadcast(self, event: dict[str, Any]) -> None:
        """Send an SSE event dict to all connected clients.

        Args:
            event: The event dictionary to broadcast.
        """
        with self._lock:
            for cq in self._clients:
                cq.put(event)

    def add_client(self) -> queue.Queue[dict[str, Any]]:
        """Register a new SSE client and return its event queue.

        Returns:
            queue.Queue[dict[str, Any]]: A queue that will receive broadcast events.
        """
        cq: queue.Queue[dict[str, Any]] = queue.Queue()
        with self._lock:
            self._clients.append(cq)
        return cq

    def remove_client(self, cq: queue.Queue[dict[str, Any]]) -> None:
        """Unregister an SSE client's event queue.

        Args:
            cq: The client queue to remove.
        """
        with self._lock:
            try:
                self._clients.remove(cq)
            except ValueError:
                logger.debug("Exception caught", exc_info=True)
                pass

    def has_clients(self) -> bool:
        with self._lock:
            return bool(self._clients)

    def _broadcast_result(
        self,
        text: str,
        step_count: int = 0,
        total_tokens: int = 0,
        cost: str = "N/A",
    ) -> None:
        event: dict[str, Any] = {
            "type": "result",
            "text": text or "(no result)",
            "step_count": step_count,
            "total_tokens": total_tokens,
            "cost": cost,
        }
        parsed = self._parse_result_yaml(text) if text else None
        if parsed:
            event["success"] = parsed.get("success")
            event["summary"] = str(parsed["summary"])
        self.broadcast(event)

    def _check_stop(self) -> None:
        if self.stop_event.is_set():
            raise KeyboardInterrupt("Agent stop requested")

    def print(self, content: Any, type: str = "text", **kwargs: Any) -> str:
        """Render content by broadcasting SSE events to connected browser clients.

        Args:
            content: The content to display.
            type: Content type (e.g. "text", "prompt", "stream_event",
                "tool_call", "tool_result", "result", "usage_info", "message").
            **kwargs: Additional options such as tool_input, is_error, cost,
                step_count, total_tokens.

        Returns:
            str: Extracted text from stream events, or empty string.
        """
        self._check_stop()
        if type == "text":
            from io import StringIO

            from rich.console import Console

            buf = StringIO()
            Console(file=buf, highlight=False, width=120, no_color=True).print(content)
            text = buf.getvalue()
            if text.strip():
                self.broadcast({"type": "text_delta", "text": text})
            return ""
        if type == "prompt":
            self.broadcast({"type": "prompt", "text": str(content)})
            return ""
        if type == "stream_event":
            return self._handle_stream_event(content)
        if type == "message":
            self._handle_message(content, **kwargs)
            return ""
        if type == "usage_info":
            self.broadcast({"type": "usage_info", "text": str(content).strip()})
            return ""
        if type == "bash_stream":
            with self._bash_lock:
                self._bash_buffer.append(str(content))
                if time.monotonic() - self._bash_last_flush >= 0.1:
                    needs_flush = True
                elif self._bash_flush_timer is None:
                    self._bash_flush_timer = threading.Timer(0.1, self._flush_bash)
                    self._bash_flush_timer.daemon = True
                    self._bash_flush_timer.start()
                    needs_flush = False
                else:
                    needs_flush = False
            if needs_flush:
                self._flush_bash()
            return ""
        if type == "tool_call":
            self._flush_bash()
            self.broadcast({"type": "text_end"})
            self._format_tool_call(str(content), kwargs.get("tool_input", {}))
            return ""
        if type == "tool_result":
            self._flush_bash()
            self.broadcast(
                {
                    "type": "tool_result",
                    "content": truncate_result(str(content)),
                    "is_error": kwargs.get("is_error", False),
                }
            )
            return ""
        if type == "result":
            self.broadcast({"type": "text_end"})
            self._broadcast_result(
                str(content),
                kwargs.get("step_count", 0),
                kwargs.get("total_tokens", 0),
                kwargs.get("cost", "N/A"),
            )
            return ""
        return ""

    async def token_callback(self, token: str) -> None:
        """Broadcast a streamed token as an SSE delta event to browser clients.

        Args:
            token: The text token to broadcast.
        """
        self._check_stop()
        if token:
            delta_type = (
                "thinking_delta" if self._current_block_type == "thinking" else "text_delta"
            )
            self.broadcast({"type": delta_type, "text": token})

    def _format_tool_call(self, name: str, tool_input: dict[str, Any]) -> None:
        file_path, lang = extract_path_and_lang(tool_input)
        event: dict[str, Any] = {"type": "tool_call", "name": name}
        if file_path:
            event["path"] = file_path
            event["lang"] = lang
        if desc := tool_input.get("description"):
            event["description"] = str(desc)
        if command := tool_input.get("command"):
            event["command"] = str(command)
        if content := tool_input.get("content"):
            event["content"] = str(content)
        old_string = tool_input.get("old_string")
        new_string = tool_input.get("new_string")
        if old_string is not None:
            event["old_string"] = str(old_string)
        if new_string is not None:
            event["new_string"] = str(new_string)
        extras = extract_extras(tool_input)
        if extras:
            event["extras"] = extras
        self.broadcast(event)

    def _handle_stream_event(self, event: Any) -> str:
        evt = event.event
        evt_type = evt.get("type", "")
        text = ""

        if evt_type == "content_block_start":
            block = evt.get("content_block", {})
            block_type = block.get("type", "")
            self._current_block_type = block_type
            if block_type == "thinking":
                self.broadcast({"type": "thinking_start"})
            elif block_type == "tool_use":
                self._tool_name = block.get("name", "?")
                self._tool_json_buffer = ""

        elif evt_type == "content_block_delta":
            delta = evt.get("delta", {})
            delta_type = delta.get("type", "")
            if delta_type == "thinking_delta":
                text = delta.get("thinking", "")
            elif delta_type == "text_delta":
                text = delta.get("text", "")
            elif delta_type == "input_json_delta":
                self._tool_json_buffer += delta.get("partial_json", "")

        elif evt_type == "content_block_stop":
            block_type = self._current_block_type
            if block_type == "thinking":
                self.broadcast({"type": "thinking_end"})
            elif block_type == "tool_use":
                try:
                    tool_input = json.loads(self._tool_json_buffer)
                except (json.JSONDecodeError, ValueError):
                    logger.debug("Exception caught", exc_info=True)
                    tool_input = {"_raw": self._tool_json_buffer}
                self._format_tool_call(self._tool_name, tool_input)
            else:
                self.broadcast({"type": "text_end"})
            self._current_block_type = ""

        return text

    def _handle_message(self, message: Any, **kwargs: Any) -> None:
        if hasattr(message, "subtype") and hasattr(message, "data"):
            if message.subtype == "tool_output":
                text = message.data.get("content", "")
                if text:
                    self.broadcast({"type": "system_output", "text": text})
        elif hasattr(message, "result"):
            budget_used = kwargs.get("budget_used", 0.0)
            self._broadcast_result(
                message.result,
                kwargs.get("step_count", 0),
                kwargs.get("total_tokens_used", 0),
                f"${budget_used:.4f}" if budget_used else "N/A",
            )
        elif hasattr(message, "content"):
            for block in message.content:
                if hasattr(block, "is_error") and hasattr(block, "content"):
                    self.broadcast(
                        {
                            "type": "tool_result",
                            "content": truncate_result(str(block.content)),
                            "is_error": bool(block.is_error),
                        }
                    )
