import uuid
from typing import TYPE_CHECKING, Dict, List, Optional

from sqlalchemy import JSON, ForeignKey, Index, String
from sqlalchemy.orm import Mapped, mapped_column, relationship

from letta.orm.mixins import ProjectMixin
from letta.orm.sqlalchemy_base import SqlalchemyBase
from letta.schemas.enums import StepStatus
from letta.schemas.step import Step as PydanticStep

if TYPE_CHECKING:
    from letta.orm.message import Message
    from letta.orm.organization import Organization
    from letta.orm.provider import Provider
    from letta.orm.run import Run
    from letta.orm.step_metrics import StepMetrics


class Step(SqlalchemyBase, ProjectMixin):
    """Tracks all metadata for agent step."""

    __tablename__ = "steps"
    __pydantic_model__ = PydanticStep
    __table_args__ = (Index("ix_steps_run_id", "run_id"),)

    id: Mapped[str] = mapped_column(String, primary_key=True, default=lambda: f"step-{uuid.uuid4()}")
    origin: Mapped[Optional[str]] = mapped_column(nullable=True, doc="The surface that this agent step was initiated from.")
    organization_id: Mapped[str] = mapped_column(
        ForeignKey("organizations.id", ondelete="RESTRICT"),
        nullable=True,
        doc="The unique identifier of the organization that this step ran for",
    )
    provider_id: Mapped[Optional[str]] = mapped_column(
        ForeignKey("providers.id", ondelete="RESTRICT"),
        nullable=True,
        doc="The unique identifier of the provider that was configured for this step",
    )
    run_id: Mapped[Optional[str]] = mapped_column(
        ForeignKey("runs.id", ondelete="SET NULL"), nullable=True, doc="The unique identifier of the run that this step belongs to"
    )
    agent_id: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the model used for this step.")
    provider_name: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the provider used for this step.")
    provider_category: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The category of the provider used for this step.")
    model: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The name of the model used for this step.")
    model_endpoint: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The model endpoint url used for this step.")
    context_window_limit: Mapped[Optional[int]] = mapped_column(
        None, nullable=True, doc="The context window limit configured for this step."
    )
    completion_tokens: Mapped[int] = mapped_column(default=0, doc="Number of tokens generated by the agent")
    prompt_tokens: Mapped[int] = mapped_column(default=0, doc="Number of tokens in the prompt")
    total_tokens: Mapped[int] = mapped_column(default=0, doc="Total number of tokens processed by the agent")
    completion_tokens_details: Mapped[Optional[Dict]] = mapped_column(
        JSON, nullable=True, doc="Detailed completion token breakdown (e.g., reasoning_tokens)."
    )
    prompt_tokens_details: Mapped[Optional[Dict]] = mapped_column(
        JSON, nullable=True, doc="Detailed prompt token breakdown (e.g., cached_tokens, cache_read_tokens, cache_creation_tokens)."
    )
    stop_reason: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The stop reason associated with this step.")
    tags: Mapped[Optional[List]] = mapped_column(JSON, doc="Metadata tags.")
    tid: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="Transaction ID that processed the step.")
    trace_id: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The trace id of the agent step.")
    request_id: Mapped[Optional[str]] = mapped_column(
        None, nullable=True, doc="The API request log ID from cloud-api for correlating steps with API requests."
    )
    feedback: Mapped[Optional[str]] = mapped_column(
        None, nullable=True, doc="The feedback for this step. Must be either 'positive' or 'negative'."
    )

    # error handling
    error_type: Mapped[Optional[str]] = mapped_column(None, nullable=True, doc="The type/class of the error that occurred")
    error_data: Mapped[Optional[Dict]] = mapped_column(
        JSON, nullable=True, doc="Error details including message, traceback, and additional context"
    )
    status: Mapped[Optional[StepStatus]] = mapped_column(None, nullable=True, doc="Step status: pending, success, or failed")

    # Relationships (foreign keys)
    organization: Mapped[Optional["Organization"]] = relationship("Organization", lazy="raise")
    provider: Mapped[Optional["Provider"]] = relationship("Provider", lazy="raise")
    run: Mapped[Optional["Run"]] = relationship("Run", back_populates="steps", lazy="raise")

    # Relationships (backrefs)
    messages: Mapped[List["Message"]] = relationship("Message", back_populates="step", cascade="save-update", lazy="noload")
    metrics: Mapped[Optional["StepMetrics"]] = relationship(
        "StepMetrics", back_populates="step", cascade="all, delete-orphan", lazy="noload", uselist=False
    )