"""Internal LLM types: stream events, tool calls, tool results."""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Any, Literal

StreamEventType = Literal[
    "text_delta", "tool_call_start", "tool_call_delta", "tool_call_end",
    "usage", "done", "error", "reasoning_delta",
]


@dataclass
class StreamEvent:
    type: StreamEventType
    data: dict[str, Any] = field(default_factory=dict)


def repair_json(text: str) -> dict[str, Any] | None:
    """Best-effort repair of a truncated/malformed JSON object (model output).

    Strategy: strip code fences, then walk backwards over plausible cut points and close the
    open brackets. Returns None when nothing parses.
    """
    s = text.strip()
    s = re.sub(r"^```(?:json)?\s*|\s*```$", "", s).strip()
    if not s:
        return None
    try:
        v = json.loads(s)
        return v if isinstance(v, dict) else None
    except json.JSONDecodeError:
        pass
    start = s.find("{")
    if start < 0:
        return None
    s = s[start:]

    def closers_at(cut: int) -> str | None:
        stack: list[str] = []
        in_str = False
        esc = False
        for ch in s[:cut]:
            if in_str:
                if esc:
                    esc = False
                elif ch == "\\":
                    esc = True
                elif ch == '"':
                    in_str = False
                continue
            if ch == '"':
                in_str = True
            elif ch in "{[":
                stack.append("}" if ch == "{" else "]")
            elif ch in "}]":
                if stack:
                    stack.pop()
        if in_str:
            return None
        return "".join(reversed(stack))

    n = len(s)
    for cut in range(n, max(0, n - 4000), -1):
        tail = s[cut - 1] if cut > 0 else ""
        if tail not in '"}]0123456789elsu.':  # plausible end of a value
            continue
        closers = closers_at(cut)
        if closers is None:
            continue
        candidate = re.sub(r",\s*$", "", s[:cut]) + closers
        try:
            v = json.loads(candidate)
            return v if isinstance(v, dict) else None
        except json.JSONDecodeError:
            continue
    return None


@dataclass
class ToolCallReq:
    id: str
    name: str
    arguments_json: str = ""
    index: int = 0
    truncated: bool = False  # the model hit its output limit while emitting this call

    def arguments(self) -> dict[str, Any]:
        if not self.arguments_json.strip():
            return {}
        try:
            v = json.loads(self.arguments_json)
            return v if isinstance(v, dict) else {}
        except json.JSONDecodeError:
            repaired = repair_json(self.arguments_json)
            if repaired is not None:
                repaired["__repaired__"] = True
                return repaired
            return {"__invalid_json__": self.arguments_json[:2000]}


@dataclass
class Artifact:
    type: str  # xlsx | image | csv | file | quiz | sources | search
    file_id: str | None = None
    filename: str | None = None
    preview: Any = None
    url: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "type": self.type,
            "file_id": self.file_id,
            "filename": self.filename,
            "preview": self.preview,
            "url": self.url,
        }


@dataclass
class ToolResult:
    content: str  # for the model
    artifacts: list[Artifact] = field(default_factory=list)
    meta: dict[str, Any] = field(default_factory=dict)
    payload: dict[str, Any] = field(default_factory=dict)  # for the UI card
    error: bool = False

    def summary(self, limit: int = 240) -> str:
        s = self.meta.get("summary") or self.content
        return s if len(s) <= limit else s[: limit - 1] + "…"


@dataclass
class Usage:
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    model: str = ""
    provider: str = ""
    generation_id: str = ""