"""Initial migration with renamed token fields.

Revision ID: 5ffc5f0266a1
Revises:
Create Date: 2025-09-11 15:08:18.429080

"""

from collections.abc import Sequence

from alembic import op
import sqlalchemy as sa
import sqlmodel.sql.sqltypes

from llmling_agent_storage.sql_provider.models import UTCDateTime


# revision identifiers, used by Alembic.
revision: str = "5ffc5f0266a1"
down_revision: str | Sequence[str] | None = None
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    """Upgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "commandhistory",
        sa.Column("id", sa.Integer(), nullable=False),
        sa.Column("session_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("agent_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("command", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("context_type", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("context_metadata", sa.JSON(), nullable=True),
        sa.Column("timestamp", UTCDateTime(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_commandhistory_agent_name"),
        "commandhistory",
        ["agent_name"],
        unique=False,
    )
    op.create_index(
        op.f("ix_commandhistory_context_type"),
        "commandhistory",
        ["context_type"],
        unique=False,
    )
    op.create_index(
        op.f("ix_commandhistory_session_id"),
        "commandhistory",
        ["session_id"],
        unique=False,
    )
    op.create_table(
        "conversation",
        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("agent_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("start_time", UTCDateTime(), nullable=True),
        sa.Column("total_tokens", sa.Integer(), nullable=False),
        sa.Column("total_cost", sa.Float(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_conversation_agent_name"), "conversation", ["agent_name"], unique=False
    )
    op.create_index(
        op.f("ix_conversation_start_time"), "conversation", ["start_time"], unique=False
    )
    op.create_table(
        "message",
        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("conversation_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("timestamp", UTCDateTime(), nullable=True),
        sa.Column("role", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("content", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("model", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("model_name", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("model_provider", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("forwarded_from", sa.JSON(), nullable=True),
        sa.Column("total_tokens", sa.Integer(), nullable=True),
        sa.Column("input_tokens", sa.Integer(), nullable=True),
        sa.Column("output_tokens", sa.Integer(), nullable=True),
        sa.Column("cost", sa.Float(), nullable=True),
        sa.Column("response_time", sa.Float(), nullable=True),
        sa.Column("checkpoint_data", sa.JSON(), nullable=True),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_message_conversation_id"), "message", ["conversation_id"], unique=False
    )
    op.create_index(op.f("ix_message_cost"), "message", ["cost"], unique=False)
    op.create_index(
        op.f("ix_message_model_name"), "message", ["model_name"], unique=False
    )
    op.create_index(
        op.f("ix_message_model_provider"), "message", ["model_provider"], unique=False
    )
    op.create_index(op.f("ix_message_name"), "message", ["name"], unique=False)
    op.create_index(
        op.f("ix_message_total_tokens"), "message", ["total_tokens"], unique=False
    )
    op.create_table(
        "toolcall",
        sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("conversation_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("message_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("timestamp", UTCDateTime(), nullable=True),
        sa.Column("tool_call_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
        sa.Column("tool_name", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.Column("args", sa.JSON(), nullable=True),
        sa.Column("result", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(
        op.f("ix_toolcall_conversation_id"), "toolcall", ["conversation_id"], unique=False
    )
    op.create_index(
        op.f("ix_toolcall_message_id"), "toolcall", ["message_id"], unique=False
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    """Downgrade schema."""
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(op.f("ix_toolcall_message_id"), table_name="toolcall")
    op.drop_index(op.f("ix_toolcall_conversation_id"), table_name="toolcall")
    op.drop_table("toolcall")
    op.drop_index(op.f("ix_message_total_tokens"), table_name="message")
    op.drop_index(op.f("ix_message_name"), table_name="message")
    op.drop_index(op.f("ix_message_model_provider"), table_name="message")
    op.drop_index(op.f("ix_message_model_name"), table_name="message")
    op.drop_index(op.f("ix_message_cost"), table_name="message")
    op.drop_index(op.f("ix_message_conversation_id"), table_name="message")
    op.drop_table("message")
    op.drop_index(op.f("ix_conversation_start_time"), table_name="conversation")
    op.drop_index(op.f("ix_conversation_agent_name"), table_name="conversation")
    op.drop_table("conversation")
    op.drop_index(op.f("ix_commandhistory_session_id"), table_name="commandhistory")
    op.drop_index(op.f("ix_commandhistory_context_type"), table_name="commandhistory")
    op.drop_index(op.f("ix_commandhistory_agent_name"), table_name="commandhistory")
    op.drop_table("commandhistory")
    # ### end Alembic commands ###
