from __future__ import annotations

import json

from cybervisor.stream_logging import CanonicalLogEvent, canonicalize_message_text

_TOOL_NAME_MAP = {
    "read_file": "Read",
    "ReadFile": "Read",
    "write_file": "Write",
    "WriteFile": "Write",
    "edit": "Edit",
    "Edit": "Edit",
    "glob": "Glob",
    "FindFiles": "Glob",
    "grep": "Grep",
    "Grep": "Grep",
    "shell": "Bash",
    "Shell": "Bash",
    "write_todos": "TodoWrite",
    "WriteTodos": "TodoWrite",
}


def _normalize_tool_name(tool_name: object) -> str | None:
    if not isinstance(tool_name, str) or not tool_name.strip():
        return None
    normalized = tool_name.strip()
    return _TOOL_NAME_MAP.get(normalized, normalized)


class GeminiStreamParser:
    def __init__(self) -> None:
        self._pending_assistant_text: list[str] = []

    def reset(self) -> None:
        self._pending_assistant_text = []

    def flush_pending_events(self) -> list[CanonicalLogEvent]:
        if not self._pending_assistant_text:
            return []
        combined_text = "".join(self._pending_assistant_text)
        self._pending_assistant_text = []
        return canonicalize_message_text(combined_text, as_hook_feedback=False)

    def parse_output_line(self, line: str) -> list[CanonicalLogEvent]:
        if not line:
            return []

        try:
            payload = json.loads(line)
        except json.JSONDecodeError:
            return [*self.flush_pending_events(), CanonicalLogEvent(kind="plain_text", text=line)]

        if not isinstance(payload, dict):
            return [*self.flush_pending_events(), CanonicalLogEvent(kind="plain_text", text=line)]

        payload_type = payload.get("type")
        if payload_type == "init":
            parts: list[str] = ["init"]
            model = payload.get("model")
            if isinstance(model, str) and model:
                parts.append(f"model={model}")
            return [
                *self.flush_pending_events(),
                CanonicalLogEvent(kind="session_start", summary_parts=tuple(parts)),
            ]

        if payload_type == "message":
            role = payload.get("role")
            content = payload.get("content")
            if role == "assistant" and isinstance(content, str) and content.strip():
                self._pending_assistant_text.append(content)
                return []
            return self.flush_pending_events()

        if payload_type == "tool_use":
            parameters = payload.get("parameters")
            return [
                *self.flush_pending_events(),
                CanonicalLogEvent(
                    kind="tool_call",
                    tool_name=_normalize_tool_name(payload.get("tool_name")),
                    tool_input=parameters if isinstance(parameters, dict) else None,
                ),
            ]

        if payload_type == "tool_result":
            return [
                *self.flush_pending_events(),
                CanonicalLogEvent(
                    kind="tool_result",
                    status=payload.get("status") if isinstance(payload.get("status"), str) else None,
                ),
            ]

        if payload_type == "error":
            message = payload.get("message")
            severity = payload.get("severity")
            if isinstance(message, str) and message.strip():
                prefix = f"{severity}: " if isinstance(severity, str) and severity.strip() else ""
                return [
                    *self.flush_pending_events(),
                    CanonicalLogEvent(kind="plain_text", text=f"{prefix}{message.strip()}"),
                ]
            return self.flush_pending_events()

        if payload_type == "result":
            summary_parts: list[str] = []
            status = payload.get("status")
            if isinstance(status, str) and status:
                summary_parts.append(status)
            stats = payload.get("stats")
            if isinstance(stats, dict):
                duration = stats.get("duration_ms")
                if isinstance(duration, (int, float)):
                    summary_parts.append(f"duration={duration}ms")
                tool_calls = stats.get("tool_calls")
                if isinstance(tool_calls, int):
                    summary_parts.append(f"tool_calls={tool_calls}")
            return [
                *self.flush_pending_events(),
                CanonicalLogEvent(
                    kind="completed",
                    summary_parts=tuple(summary_parts),
                    status=status if isinstance(status, str) else None,
                ),
            ]

        return [*self.flush_pending_events(), CanonicalLogEvent(kind="plain_text", text=line)]
