"""Traceability payload schemas (KM-691).

User-facing provenance for one assistant turn: what the AI planned, which tools it
called (with real inputs + outputs), and which data sources it read (with the
executed query). One `TraceabilityPayload` is built per assistant `message_id`,
stored as a JSONB row, and served by `GET /api/v1/traceability`.

Distinct from Langfuse *observability* (`src/observability/langfuse/`): that is
engineering-only, PII-masked (arg keys + row counts). Traceability shows the user
their own data's provenance — real args, output previews, executed SQL — so it is
NOT masked. Truncation caps (in `scratchpad.py`) bound the payload size instead.

`thinking` is always `null` in v1 (our agents are plain chat completions with no
native reasoning output; synthesizing it post-hoc would be unfaithful). The field
stays in the payload so adding it later is contract-compatible.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field


class PlanStep(BaseModel):
    """One CRISP-DM step the Planner ran (derived from an AnalysisRecord task)."""

    step: int
    stage: str
    objective: str
    status: str
    tools_used: list[str] = Field(default_factory=list)


class PlanningInfo(BaseModel):
    """Planner output for a slow-path turn; `null` on every other turn type."""

    goal_restated: str
    assumptions: list[str] = Field(default_factory=list)
    steps: list[PlanStep] = Field(default_factory=list)


class ToolCallInfo(BaseModel):
    """One tool invocation with its real input + output (both truncation-capped).

    The field is spelled `input` in the wire contract (the FE reads it), which
    shadows the `input` builtin — hence the alias + `populate_by_name` so callers
    can pass `input=` on construction and `model_dump(by_alias=True)` emits `input`.
    """

    model_config = ConfigDict(populate_by_name=True)

    order: int
    task_id: str | None = None
    name: str
    input_: dict[str, Any] = Field(default_factory=dict, alias="input")
    output: dict[str, Any] = Field(default_factory=dict)
    status: Literal["success", "error"] = "success"
    error: str | None = None


class TraceabilityPayload(BaseModel):
    """The full provenance record for one assistant `message_id`."""

    model_config = ConfigDict(populate_by_name=True)

    analysis_id: str
    message_id: str
    user_id: str  # ownership column on the row; the FE may ignore it
    intent: str
    generated_at: datetime
    planning: PlanningInfo | None = None
    thinking: str | None = None
    tool_calls: list[ToolCallInfo] = Field(default_factory=list)
    sources: list[dict[str, Any]] = Field(default_factory=list)