"""Beaver session 子系统的数据模型。

这层只定义数据结构，不放数据库读写逻辑。目的是把：
1. SQLite 行结构
2. 运行时会话对象
3. 对外暴露的 conversation message

三件事分开，避免后续所有地方都直接和裸字典耦合。
"""

from __future__ import annotations

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


@dataclass(slots=True)
class SessionUsage:
    """会话维度的 usage/cost 统计。"""

    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    cache_write_tokens: int = 0
    reasoning_tokens: int = 0
    estimated_cost_usd: float = 0.0
    actual_cost_usd: float | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cache_read_tokens": self.cache_read_tokens,
            "cache_write_tokens": self.cache_write_tokens,
            "reasoning_tokens": self.reasoning_tokens,
            "estimated_cost_usd": self.estimated_cost_usd,
            "actual_cost_usd": self.actual_cost_usd,
        }


@dataclass(slots=True)
class MessageRecord:
    """单条会话事件的结构化表示。

    当前仍然沿用 `messages` 这张表名，但语义已经开始向 event stream 收拢：
    1. 普通 user/assistant/tool 消息本身就是事件
    2. 运行时的 system snapshot / run lifecycle 也可写成隐藏事件
    3. 是否进入模型上下文由 `context_visible` 决定，而不是简单看 role
    """

    role: str
    content: str | None = None
    timestamp: float | None = None
    message_id: int | None = None
    run_id: str | None = None
    event_type: str | None = None
    event_payload: dict[str, Any] | None = None
    context_visible: bool = True
    tool_name: str | None = None
    tool_calls: list[dict[str, Any]] | None = None
    tool_call_id: str | None = None
    finish_reason: str | None = None
    reasoning: str | None = None
    reasoning_details: Any | None = None
    codex_reasoning_items: Any | None = None

    def to_conversation_message(self) -> dict[str, Any]:
        """转成 provider / context builder 可直接消费的消息格式。"""

        if not self.context_visible:
            raise ValueError("Hidden session events cannot be converted into conversation messages")

        payload: dict[str, Any] = {
            "role": self.role,
            "content": self.content,
        }
        if self.timestamp is not None:
            payload["timestamp"] = self.timestamp
        if self.run_id:
            payload["run_id"] = self.run_id
        if self.event_payload:
            if self.event_payload.get("task_id"):
                payload["task_id"] = self.event_payload.get("task_id")
            if self.event_payload.get("task_status"):
                payload["task_status"] = self.event_payload.get("task_status")
            if self.event_payload.get("evidence_status"):
                payload["evidence_status"] = self.event_payload.get("evidence_status")
            if self.event_payload.get("acceptance_state"):
                payload["acceptance_state"] = self.event_payload.get("acceptance_state")
            if self.event_payload.get("feedback_state"):
                payload["feedback_state"] = self.event_payload.get("feedback_state")
            if self.event_payload.get("feedback_error"):
                payload["feedback_error"] = self.event_payload.get("feedback_error")
            for key in (
                "message_type",
                "scheduled_job_id",
                "scheduled_run_id",
                "cron_job_name",
                "mode",
            ):
                if self.event_payload.get(key):
                    payload[key] = self.event_payload.get(key)
        if self.tool_name:
            payload["tool_name"] = self.tool_name
        if self.tool_calls:
            payload["tool_calls"] = self.tool_calls
        if self.tool_call_id:
            payload["tool_call_id"] = self.tool_call_id
        if self.finish_reason:
            payload["finish_reason"] = self.finish_reason
        if self.reasoning:
            payload["reasoning"] = self.reasoning
        if self.reasoning_details is not None:
            payload["reasoning_details"] = self.reasoning_details
        if self.codex_reasoning_items is not None:
            payload["codex_reasoning_items"] = self.codex_reasoning_items
        return payload

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> "MessageRecord":
        """从 SQLite row/dict 恢复消息模型。"""

        tool_calls = row.get("tool_calls")
        if isinstance(tool_calls, str):
            try:
                tool_calls = json.loads(tool_calls)
            except json.JSONDecodeError:
                tool_calls = []

        reasoning_details = row.get("reasoning_details")
        if isinstance(reasoning_details, str):
            try:
                reasoning_details = json.loads(reasoning_details)
            except json.JSONDecodeError:
                reasoning_details = None

        codex_reasoning_items = row.get("codex_reasoning_items")
        if isinstance(codex_reasoning_items, str):
            try:
                codex_reasoning_items = json.loads(codex_reasoning_items)
            except json.JSONDecodeError:
                codex_reasoning_items = None

        event_payload = row.get("event_payload")
        if isinstance(event_payload, str):
            try:
                event_payload = json.loads(event_payload)
            except json.JSONDecodeError:
                event_payload = None

        return cls(
            message_id=row.get("id"),
            run_id=row.get("run_id"),
            role=row["role"],
            content=row.get("content"),
            event_type=row.get("event_type") or row.get("role"),
            event_payload=event_payload,
            context_visible=bool(row.get("context_visible", 1)),
            tool_name=row.get("tool_name"),
            tool_calls=tool_calls,
            tool_call_id=row.get("tool_call_id"),
            timestamp=row.get("timestamp"),
            finish_reason=row.get("finish_reason"),
            reasoning=row.get("reasoning"),
            reasoning_details=reasoning_details,
            codex_reasoning_items=codex_reasoning_items,
        )


@dataclass(slots=True)
class SessionRecord:
    """单个 session 的结构化表示。"""

    session_id: str
    source: str
    started_at: float
    last_active: float
    user_id: str | None = None
    title: str | None = None
    model: str | None = None
    system_prompt: str | None = None
    parent_session_id: str | None = None
    ended_at: float | None = None
    end_reason: str | None = None
    message_count: int = 0
    tool_call_count: int = 0
    preview: str | None = None
    usage: SessionUsage = field(default_factory=SessionUsage)

    def to_dict(self) -> dict[str, Any]:
        payload = {
            "id": self.session_id,
            "source": self.source,
            "user_id": self.user_id,
            "title": self.title,
            "model": self.model,
            "system_prompt": self.system_prompt,
            "parent_session_id": self.parent_session_id,
            "started_at": self.started_at,
            "last_active": self.last_active,
            "ended_at": self.ended_at,
            "end_reason": self.end_reason,
            "message_count": self.message_count,
            "tool_call_count": self.tool_call_count,
            "preview": self.preview,
        }
        payload.update(self.usage.to_dict())
        return payload

    @classmethod
    def from_row(cls, row: dict[str, Any]) -> "SessionRecord":
        return cls(
            session_id=row["id"],
            source=row["source"],
            user_id=row.get("user_id"),
            title=row.get("title"),
            model=row.get("model"),
            system_prompt=row.get("system_prompt"),
            parent_session_id=row.get("parent_session_id"),
            started_at=row["started_at"],
            last_active=row["last_active"],
            ended_at=row.get("ended_at"),
            end_reason=row.get("end_reason"),
            message_count=row.get("message_count", 0),
            tool_call_count=row.get("tool_call_count", 0),
            preview=row.get("preview"),
            usage=SessionUsage(
                input_tokens=row.get("input_tokens", 0),
                output_tokens=row.get("output_tokens", 0),
                cache_read_tokens=row.get("cache_read_tokens", 0),
                cache_write_tokens=row.get("cache_write_tokens", 0),
                reasoning_tokens=row.get("reasoning_tokens", 0),
                estimated_cost_usd=row.get("estimated_cost_usd", 0.0) or 0.0,
                actual_cost_usd=row.get("actual_cost_usd"),
            ),
        )
