"""SQLAlchemy ORM models for Run and Step.

This module defines the database schema for storing pipeline execution traces.
Uses PostgreSQL with JSONB for flexible JSON storage with indexing capability.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any
from uuid import UUID

from sqlalchemy import JSON, DateTime, ForeignKey, Index, String, Text, text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from sqlalchemy.types import TypeDecorator


class JSONType(TypeDecorator):
    """JSON type that uses JSONB on PostgreSQL and JSON on other databases.

    This allows tests to run on SQLite while production uses PostgreSQL JSONB.
    """

    impl = JSON
    cache_ok = True

    def load_dialect_impl(self, dialect):
        if dialect.name == "postgresql":
            return dialect.type_descriptor(JSONB())
        return dialect.type_descriptor(JSON())


class Base(DeclarativeBase):
    """Base class for all ORM models."""

    pass


class Run(Base):
    """Run model - represents a complete pipeline execution.

    A Run captures the full lifecycle of a pipeline invocation, from start
    to completion (success or error). It contains metadata about the pipeline,
    input/output summaries, and links to all Steps executed within the run.

    Attributes:
        id: Unique identifier (UUID generated by SDK)
        pipeline_name: Name of the pipeline (e.g., "recommendation_pipeline")
        status: Current status ("running", "success", "error")
        started_at: When the run started
        ended_at: When the run completed (None if still running)
        input_summary: Summarized input data (JSONB)
        output_summary: Summarized output data (JSONB)
        metadata_: Additional metadata (request_id, user_id, etc.)
        request_id: Extracted for indexing (from metadata)
        user_id: Extracted for indexing (from metadata)
        environment: Extracted for indexing (from metadata)
        error_message: Error details if status is "error"
        steps: Related Step records (ordered by index)
    """

    __tablename__ = "runs"

    id: Mapped[UUID] = mapped_column(primary_key=True)
    pipeline_name: Mapped[str] = mapped_column(String(255))
    status: Mapped[str] = mapped_column(String(20), default="running")
    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
    ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None)
    input_summary: Mapped[dict[str, Any] | None] = mapped_column(JSONType, default=None)
    output_summary: Mapped[dict[str, Any] | None] = mapped_column(JSONType, default=None)
    metadata_: Mapped[dict[str, Any] | None] = mapped_column(
        "metadata", JSONType, default=None
    )
    request_id: Mapped[str | None] = mapped_column(String(255), default=None)
    user_id: Mapped[str | None] = mapped_column(String(255), default=None)
    environment: Mapped[str | None] = mapped_column(String(50), default=None)
    error_message: Mapped[str | None] = mapped_column(Text, default=None)

    # Relationships
    steps: Mapped[list[Step]] = relationship(
        back_populates="run",
        order_by="Step.index",
        cascade="all, delete-orphan",
    )

    __table_args__ = (
        Index("ix_runs_pipeline_name", "pipeline_name"),
        Index("ix_runs_status", "status"),
        Index("ix_runs_started_at", "started_at"),
        Index("ix_runs_request_id", "request_id"),
        Index("ix_runs_user_id", "user_id"),
        Index("ix_runs_environment", "environment"),
    )


class Step(Base):
    """Step model - represents a single step within a pipeline run.

    A Step captures one processing stage (filter, rank, LLM call, etc.)
    including its inputs, outputs, timing, and reasoning.

    Attributes:
        id: Unique identifier (UUID generated by SDK)
        run_id: Foreign key to parent Run
        step_name: Human-readable name (e.g., "price_filter")
        step_type: Type of step ("filter", "rank", "llm", etc.)
        index: Order within the run (0-based)
        started_at: When the step started
        ended_at: When the step completed
        duration_ms: Execution time in milliseconds
        input_summary: Summarized input data (JSONB)
        output_summary: Summarized output data (JSONB)
        input_count: Number of input items (for list-like inputs)
        output_count: Number of output items (for list-like outputs)
        reasoning: Why decisions were made (JSONB)
        metadata_: Additional step-specific metadata (JSONB)
        status: Step status ("running", "success", "error")
        error_message: Error details if status is "error"
        run: Parent Run relationship
    """

    __tablename__ = "steps"

    id: Mapped[UUID] = mapped_column(primary_key=True)
    run_id: Mapped[UUID] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"))
    step_name: Mapped[str] = mapped_column(String(255))
    step_type: Mapped[str] = mapped_column(String(50))
    index: Mapped[int]
    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True))
    ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), default=None)
    duration_ms: Mapped[int | None] = mapped_column(default=None)
    input_summary: Mapped[dict[str, Any] | None] = mapped_column(JSONType, default=None)
    output_summary: Mapped[dict[str, Any] | None] = mapped_column(JSONType, default=None)
    input_count: Mapped[int | None] = mapped_column(default=None)
    output_count: Mapped[int | None] = mapped_column(default=None)
    reasoning: Mapped[dict[str, Any] | None] = mapped_column(JSONType, default=None)
    metadata_: Mapped[dict[str, Any] | None] = mapped_column(
        "metadata", JSONType, default=None
    )
    status: Mapped[str | None] = mapped_column(String(20), default=None)
    error_message: Mapped[str | None] = mapped_column(Text, default=None)

    # Relationships
    run: Mapped[Run] = relationship(back_populates="steps")

    __table_args__ = (
        Index("ix_steps_run_id", "run_id"),
        Index("ix_steps_step_name", "step_name"),
        Index("ix_steps_step_type", "step_type"),
        Index("ix_steps_input_count", "input_count"),
        Index("ix_steps_output_count", "output_count"),
        Index("ix_steps_status", "status"),
    )


class Payload(Base):
    """Payload model - stores externalized large data.

    When the SDK encounters large data (lists >100 items, strings >1KB),
    it extracts them into `_payloads` and stores references (`_ref`) in
    the summaries. This table stores those externalized payloads separately
    for efficient storage and retrieval.

    The summaries contain references like:
        {"_type": "list", "_count": 500, "_ref": "p-001", "_preview": [...]}

    The actual data for "p-001" is stored here.

    Attributes:
        id: Auto-generated primary key
        run_id: Parent Run UUID
        step_id: Parent Step UUID (nullable - run-level payloads have no step)
        ref_id: Reference ID from SDK (e.g., "p-001")
        phase: "input" or "output" - which phase the payload came from
        data: The actual payload data (can be any JSON-serializable value)
    """

    __tablename__ = "payloads"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    run_id: Mapped[UUID] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE"))
    step_id: Mapped[UUID | None] = mapped_column(
        ForeignKey("steps.id", ondelete="CASCADE"), default=None
    )
    ref_id: Mapped[str] = mapped_column(String(50))
    phase: Mapped[str] = mapped_column(String(20))  # "input" or "output"
    data: Mapped[Any] = mapped_column(JSONType)  # The actual payload data

    __table_args__ = (
        Index("ix_payloads_run_id", "run_id"),
        Index("ix_payloads_step_id", "step_id"),
        Index("ix_payloads_ref_id", "ref_id"),
        # Two partial indexes for uniqueness - handles NULL step_id correctly
        # (NULL != NULL in SQL, so a single composite index won't enforce uniqueness)
        # Note: postgresql_where only works on PostgreSQL. SQLite tests won't
        # enforce partial uniqueness, but production PostgreSQL will.
        # 1. For step-level payloads (step_id IS NOT NULL)
        Index(
            "ix_payloads_step_lookup",
            "run_id",
            "step_id",
            "ref_id",
            "phase",
            unique=True,
            postgresql_where=text("step_id IS NOT NULL"),
        ),
        # 2. For run-level payloads (step_id IS NULL)
        Index(
            "ix_payloads_run_lookup",
            "run_id",
            "ref_id",
            "phase",
            unique=True,
            postgresql_where=text("step_id IS NULL"),
        ),
    )
