import json
import logging
from typing import AsyncGenerator

logger = logging.getLogger(__name__)


async def parse_deerflow_sse_stream(response) -> AsyncGenerator[dict, None]:
    """
    Parse DeferFlow /runs/stream SSE response into standardized events.

    Handles LangGraph Server SSE format with event: and data: lines.
    Supports stream_mode "values" and "messages".

    Yields dicts: {"type": str, "content": str, "name": str, "arguments": str}
    """
    current_event = ""
    thinking_started = False
    prev_thinking = ""
    prev_content = ""
    seen_tool_calls = set()  # dedup by (name, arguments)

    buffer = ""
    async for chunk in response.aiter_text():
        buffer += chunk
        import re
        while "\n" in buffer:
            line, buffer = buffer.split("\n", 1)
            line = line.strip()

            if not line:
                current_event = ""
                continue

            if line.startswith("event:"):
                current_event = line[6:].strip()
                if current_event == "end":
                    yield {"type": "done"}
                    return
                continue

            if not line.startswith("data:"):
                continue

            data_str = line[5:].strip()
            if not data_str:
                continue

            try:
                parsed = json.loads(data_str)
            except json.JSONDecodeError:
                continue

            # --- Handle stream_mode: "values" ---
            if current_event == "values":
                if not isinstance(parsed, dict):
                    continue
                if "messages" in parsed and isinstance(parsed["messages"], list):
                    for msg in parsed["messages"]:
                        if not isinstance(msg, dict):
                            continue
                        for event in _extract_message_events(msg, prev_thinking, prev_content, seen_tool_calls):
                            yield event
                            if event["type"] == "thinking":
                                prev_thinking += event.get("content", "")
                            if event["type"] == "content":
                                prev_content += event.get("content", "")
                elif "content" in parsed:
                    content = str(parsed["content"])
                    if content:
                        yield {"type": "content", "content": content}

            # --- Handle stream_mode: "messages" ---
            elif current_event in ("messages/partial", "messages/complete"):
                if isinstance(parsed, list):
                    for msg in parsed:
                        if not isinstance(msg, dict):
                            continue
                        for event in _extract_message_events(msg, prev_thinking, prev_content, seen_tool_calls):
                            yield event
                            if event["type"] == "thinking":
                                prev_thinking += event.get("content", "")
                            if event["type"] == "content":
                                prev_content += event.get("content", "")
                elif isinstance(parsed, dict):
                    for event in _extract_message_events(parsed, prev_thinking, prev_content, seen_tool_calls):
                        yield event
                        if event["type"] == "thinking":
                            prev_thinking += event.get("content", "")
                        if event["type"] == "content":
                            prev_content += event.get("content", "")

            elif current_event == "error":
                yield {"type": "error", "content": str(parsed.get("message", str(parsed)))}
                return

            current_event = ""

    yield {"type": "done"}


def _extract_message_events(msg: dict, prev_thinking: str = "", prev_content: str = "", seen_tool_calls=None):
    """Extract thinking/content/tool_call/tool_result from a single message dict."""
    if seen_tool_calls is None:
        seen_tool_calls = set()

    msg_type = msg.get("type", msg.get("role", ""))
    additional = msg.get("additional_kwargs", {}) or {}
    tool_calls = msg.get("tool_calls", []) or []

    if msg_type in ("ai", "assistant"):
        # LangChain DeepSeek stores reasoning_content in additional_kwargs["reasoning_content"]
        thinking_text = additional.get("thinking", "") or additional.get("reasoning_content", "")

        if isinstance(msg.get("content"), list):
            for part in msg.get("content", []):
                if isinstance(part, dict):
                    if part.get("type") == "thinking":
                        thinking_text += part.get("thinking", "")
                    elif part.get("type") == "reasoning":
                        # LangChain extracts reasoning_content into {"type": "reasoning", "reasoning": "..."}
                        reasoning = part.get("reasoning", "")
                        if reasoning:
                            thinking_text += reasoning
                    elif part.get("type") == "text":
                        text = part.get("text", "")
                        if text:
                            yield {"type": "content", "content": text}

        if thinking_text and not prev_thinking:
            # First time seeing thinking — yield all of it
            yield {"type": "thinking", "content": thinking_text}
        elif thinking_text and len(thinking_text) > len(prev_thinking) and thinking_text.startswith(prev_thinking):
            # Incremental delta
            delta = thinking_text[len(prev_thinking):]
            if delta:
                yield {"type": "thinking", "content": delta}

        for tc in tool_calls:
            fn = tc.get("function", {}) or {}
            tc_key = (fn.get("name", ""), fn.get("arguments", ""))
            if tc_key in seen_tool_calls:
                continue
            seen_tool_calls.add(tc_key)
            yield {
                "type": "tool_call",
                "name": fn.get("name", ""),
                "arguments": fn.get("arguments", ""),
            }

        content_text = msg.get("content", "") or ""
        if isinstance(content_text, str) and content_text != prev_content:
            import re

            # Extract <think> content as thinking events (first time only)
            if not prev_thinking:
                think_matches = list(re.finditer(r'<think([\s\S]*?)<\/think>', content_text))
                for match in think_matches:
                    think_content = match.group(1).strip()
                    if think_content:
                        yield {"type": "thinking", "content": think_content}

            cleaned = re.sub(r'<think[\s\S]*?<\/think>', '', content_text).strip()
            delta = cleaned[len(prev_content):]
            if delta:
                yield {"type": "content", "content": delta}

    elif msg_type == "tool":
        tool_name = msg.get("name", "")
        tool_content = str(msg.get("content", ""))

        yield {
            "type": "tool_result",
            "name": tool_name,
            "content": tool_content,
        }


def parse_deerflow_stream_values(response) -> AsyncGenerator[dict, None]:
    """Backward-compatible alias for parse_deerflow_sse_stream."""
    return parse_deerflow_sse_stream(response)


def format_sse_event(event_type: str, payload: dict) -> str:
    return f"data: {json.dumps({'type': event_type, **payload}, ensure_ascii=False)}\n\n"