"""
StreamingCallback：拦截 agent 的输出，转成结构化 SSE 事件

后端通过 patch agent 的 print 输出方法捕获步骤事件，并转成 SSE。
"""
import io
import queue
import re
import sys


class StreamingCallback:
    """
    拦截 agent.go() 的 stdout 输出，解析成结构化事件。
    
    Biomni 的 agent 通过 print() 输出中间步骤：
    - "Thought: ..." → thinking 事件
    - "```python\n...\n```" → code 事件  
    - "Observation: ..." → observation 事件
    - 最终答案 → result 事件
    """

    def __init__(self, event_queue: queue.Queue):
        self.q = event_queue
        self._original_stdout = None
        self._buf = io.StringIO()
        self._in_solution = False
        self._solution_lines: list[str] = []

    def attach(self, _agent):
        """patch sys.stdout 来捕获 agent 的 print 输出"""
        self._original_stdout = sys.stdout
        sys.stdout = _InterceptWriter(self._parse_and_emit, self._original_stdout)

    def detach(self, _agent):
        # 兜底：若流结束时仍在 solution 块内，仍然输出 result
        self._flush_solution_if_any()
        if self._original_stdout:
            sys.stdout = self._original_stdout

    def _parse_and_emit(self, text: str):
        """解析 agent 输出文本，转成结构化事件"""
        raw_text = text.rstrip("\n")
        if self._handle_solution_markup(raw_text):
            return

        text = raw_text.strip()
        if not text:
            return

        # 匹配 Biomni A1 的输出格式
        if text.startswith("Thought:") or text.startswith("[Thinking]"):
            self.q.put({"type": "thinking", "content": text})

        elif "```python" in text or "```r" in text or "```bash" in text:
            # 提取代码块
            match = re.search(r"```(\w+)\n(.*?)```", text, re.DOTALL)
            if match:
                lang, code = match.group(1), match.group(2)
                self.q.put({"type": "code", "lang": lang, "content": code})
            else:
                self.q.put({"type": "thinking", "content": text})

        elif text.startswith("Observation:") or text.startswith("[Output]"):
            self.q.put({"type": "observation", "content": text})

        elif text.startswith("Action:") or "tool" in text.lower()[:30]:
            # 尝试提取工具名
            tool_match = re.search(r"(\w+)\(", text)
            tool_name = tool_match.group(1) if tool_match else "tool"
            self.q.put({"type": "tool_use", "tool": tool_name, "content": text})

        elif text.startswith("Final Answer:") or text.startswith("Answer:"):
            self.q.put({"type": "result", "content": text})

        else:
            # 其他输出作为 thinking
            self.q.put({"type": "thinking", "content": text})

    def _handle_solution_markup(self, raw_text: str) -> bool:
        """
        解析 <solution>...</solution> 多行块。
        命中时输出 type=result，并返回 True 表示该行已被消费。
        """
        if self._in_solution:
            if "</solution>" in raw_text:
                before_end, _after_end = raw_text.split("</solution>", 1)
                if before_end:
                    self._solution_lines.append(before_end)
                self._flush_solution_if_any()
                self._in_solution = False
            else:
                self._solution_lines.append(raw_text)
            return True

        if "<solution>" in raw_text:
            _before_start, after_start = raw_text.split("<solution>", 1)
            self._in_solution = True
            self._solution_lines = []

            if "</solution>" in after_start:
                content, _after_end = after_start.split("</solution>", 1)
                if content:
                    self._solution_lines.append(content)
                self._flush_solution_if_any()
                self._in_solution = False
            else:
                if after_start:
                    self._solution_lines.append(after_start)
            return True

        return False

    def _flush_solution_if_any(self):
        content = "\n".join(self._solution_lines).strip()
        self._solution_lines = []
        if content:
            self.q.put({"type": "result", "content": content})


class _InterceptWriter:
    """替换 sys.stdout，拦截写入"""
    def __init__(self, callback, original):
        self._cb = callback
        self._orig = original
        self._line_buf = ""

    def write(self, text):
        self._orig.write(text)  # 同时写到真实 stdout（方便调试）
        self._line_buf += text
        if "\n" in self._line_buf:
            lines = self._line_buf.split("\n")
            for line in lines[:-1]:
                self._cb(line)
            self._line_buf = lines[-1]

    def flush(self):
        if self._line_buf:
            self._cb(self._line_buf)
            self._line_buf = ""
        self._orig.flush()

    def __getattr__(self, name):
        return getattr(self._orig, name)