"""add_token_cost_and_input_output_content

Revision ID: aa259a314eeb
Revises: 83d01e7ad41c
Create Date: 2025-11-03 10:39:04.544418

"""

import sys
from pathlib import Path

import sqlalchemy as sa
from sqlalchemy import text

from alembic import op

# Add the src directory to path for imports
migration_dir = Path(__file__).parent.parent.parent
src_dir = migration_dir / "src"
sys.path.insert(0, str(src_dir))

from utils.trace import extract_token_cost_from_span, value_to_string

# revision identifiers, used by Alembic.
revision = "aa259a314eeb"
down_revision = "83d01e7ad41c"
branch_labels = None
depends_on = None

BATCH_SIZE = 50
TRACE_BATCH_SIZE = 10  # Smaller batch for trace metadata (loads many spans per trace)


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column("spans", sa.Column("prompt_token_count", sa.Integer(), nullable=True))
    op.add_column(
        "spans",
        sa.Column("completion_token_count", sa.Integer(), nullable=True),
    )
    op.add_column("spans", sa.Column("total_token_count", sa.Integer(), nullable=True))
    op.add_column("spans", sa.Column("prompt_token_cost", sa.Float(), nullable=True))
    op.add_column(
        "spans",
        sa.Column("completion_token_cost", sa.Float(), nullable=True),
    )
    op.add_column("spans", sa.Column("total_token_cost", sa.Float(), nullable=True))
    op.create_index(
        "idx_spans_total_token_cost",
        "spans",
        ["total_token_cost"],
        unique=False,
    )
    op.create_index(
        "idx_spans_total_token_count",
        "spans",
        ["total_token_count"],
        unique=False,
    )
    op.add_column(
        "trace_metadata",
        sa.Column("prompt_token_count", sa.Integer(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("completion_token_count", sa.Integer(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("total_token_count", sa.Integer(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("prompt_token_cost", sa.Float(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("completion_token_cost", sa.Float(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("total_token_cost", sa.Float(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("input_content", sa.String(), nullable=True),
    )
    op.add_column(
        "trace_metadata",
        sa.Column("output_content", sa.String(), nullable=True),
    )
    op.create_index(
        "idx_traces_total_token_cost",
        "trace_metadata",
        ["total_token_cost"],
        unique=False,
    )
    op.create_index(
        "idx_traces_total_token_count",
        "trace_metadata",
        ["total_token_count"],
        unique=False,
    )
    # ### end Alembic commands ###

    _backfill_spans_token_data()
    _backfill_trace_metadata_token_data()
    _backfill_trace_metadata_input_output_content()


def _backfill_spans_token_data():
    """Backfill spans table with token counts and costs in a single batch processing pass."""
    connection = op.get_bind()

    # Get total count for progress tracking (only LLM spans)
    result = connection.execute(
        text("SELECT COUNT(*) FROM spans WHERE span_kind = 'LLM'"),
    )
    total_spans = result.scalar()

    if total_spans == 0:
        return

    offset = 0

    while True:
        # Fetch batch of LLM spans only
        result = connection.execute(
            text(
                """
                SELECT id, raw_data
                FROM spans
                WHERE span_kind = 'LLM'
                ORDER BY created_at
                LIMIT :limit OFFSET :offset
            """,
            ),
            {"limit": BATCH_SIZE, "offset": offset},
        )

        batch = result.fetchall()
        if not batch:
            break

        # Process each span in the batch
        for span_id, raw_data in batch:
            try:
                # Use existing utility to extract/compute token counts and costs
                token_data = extract_token_cost_from_span(raw_data, "LLM")

                # Update the span with all computed values
                connection.execute(
                    text(
                        """
                        UPDATE spans
                        SET prompt_token_count = :prompt_tokens,
                            completion_token_count = :completion_tokens,
                            total_token_count = :total_tokens,
                            prompt_token_cost = :prompt_cost,
                            completion_token_cost = :completion_cost,
                            total_token_cost = :total_cost
                        WHERE id = :span_id
                    """,
                    ),
                    {
                        "prompt_tokens": token_data.prompt_token_count,
                        "completion_tokens": token_data.completion_token_count,
                        "total_tokens": token_data.total_token_count,
                        "prompt_cost": token_data.prompt_token_cost,
                        "completion_cost": token_data.completion_token_cost,
                        "total_cost": token_data.total_token_cost,
                        "span_id": span_id,
                    },
                )

            except Exception:
                # Skip spans that fail to process
                pass

            offset += 1


def _backfill_trace_metadata_token_data():
    """Backfill trace_metadata by aggregating token counts and costs from spans in batches."""
    connection = op.get_bind()

    # Get total count of traces
    result = connection.execute(text("SELECT COUNT(*) FROM trace_metadata"))
    total_traces = result.scalar()

    if total_traces == 0:
        return

    offset = 0

    while True:
        # Fetch batch of trace IDs (smaller batch size due to span aggregation)
        result = connection.execute(
            text(
                """
                SELECT trace_id
                FROM trace_metadata
                ORDER BY created_at
                LIMIT :limit OFFSET :offset
            """,
            ),
            {"limit": TRACE_BATCH_SIZE, "offset": offset},
        )

        batch = result.fetchall()
        if not batch:
            break

        trace_ids = [row[0] for row in batch]

        # Aggregate token data for this batch of traces
        result = connection.execute(
            text(
                """
                SELECT
                    trace_id,
                    SUM(prompt_token_count) as prompt_token_count,
                    SUM(completion_token_count) as completion_token_count,
                    SUM(total_token_count) as total_token_count,
                    SUM(prompt_token_cost) as prompt_token_cost,
                    SUM(completion_token_cost) as completion_token_cost,
                    SUM(total_token_cost) as total_token_cost
                FROM spans
                WHERE trace_id = ANY(:trace_ids)
                GROUP BY trace_id
            """,
            ),
            {"trace_ids": trace_ids},
        )

        aggregations = result.fetchall()

        # Update each trace in the batch
        for agg_row in aggregations:
            (
                trace_id,
                prompt_count,
                completion_count,
                total_count,
                prompt_cost,
                completion_cost,
                total_cost,
            ) = agg_row

            connection.execute(
                text(
                    """
                    UPDATE trace_metadata
                    SET prompt_token_count = :prompt_count,
                        completion_token_count = :completion_count,
                        total_token_count = :total_count,
                        prompt_token_cost = :prompt_cost,
                        completion_token_cost = :completion_cost,
                        total_token_cost = :total_cost
                    WHERE trace_id = :trace_id
                """,
                ),
                {
                    "prompt_count": prompt_count,
                    "completion_count": completion_count,
                    "total_count": total_count,
                    "prompt_cost": prompt_cost,
                    "completion_cost": completion_cost,
                    "total_cost": total_cost,
                    "trace_id": trace_id,
                },
            )

        offset += len(batch)


def _backfill_trace_metadata_input_output_content():
    """Backfill trace_metadata input/output content from root spans in batches."""
    connection = op.get_bind()

    # Get total count of traces
    result = connection.execute(text("SELECT COUNT(*) FROM trace_metadata"))
    total_traces = result.scalar()

    if total_traces == 0:
        return

    offset = 0

    while True:
        # Fetch batch of trace IDs (smaller batch size due to span queries)
        result = connection.execute(
            text(
                """
                SELECT trace_id
                FROM trace_metadata
                ORDER BY created_at
                LIMIT :limit OFFSET :offset
            """,
            ),
            {"limit": TRACE_BATCH_SIZE, "offset": offset},
        )

        batch = result.fetchall()
        if not batch:
            break

        trace_ids = [row[0] for row in batch]

        # For each trace in the batch, find the earliest root span
        result = connection.execute(
            text(
                """
                SELECT DISTINCT ON (trace_id)
                    trace_id,
                    raw_data->'attributes'->'input'->'value' as input_value,
                    raw_data->'attributes'->'output'->'value' as output_value
                FROM spans
                WHERE parent_span_id IS NULL
                  AND trace_id = ANY(:trace_ids)
                ORDER BY trace_id, start_time ASC
            """,
            ),
            {"trace_ids": trace_ids},
        )

        root_spans = result.fetchall()

        # Update each trace with its input/output content
        for row in root_spans:
            trace_id, input_value, output_value = row

            # Convert values to strings using utility function
            input_content = value_to_string(input_value)
            output_content = value_to_string(output_value)

            # Update trace_metadata with the extracted content
            connection.execute(
                text(
                    """
                    UPDATE trace_metadata
                    SET input_content = :input_content,
                        output_content = :output_content
                    WHERE trace_id = :trace_id
                """,
                ),
                {
                    "input_content": input_content,
                    "output_content": output_content,
                    "trace_id": trace_id,
                },
            )

        offset += len(batch)


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index("idx_traces_total_token_count", table_name="trace_metadata")
    op.drop_index("idx_traces_total_token_cost", table_name="trace_metadata")
    op.drop_column("trace_metadata", "output_content")
    op.drop_column("trace_metadata", "input_content")
    op.drop_column("trace_metadata", "total_token_cost")
    op.drop_column("trace_metadata", "completion_token_cost")
    op.drop_column("trace_metadata", "prompt_token_cost")
    op.drop_column("trace_metadata", "total_token_count")
    op.drop_column("trace_metadata", "completion_token_count")
    op.drop_column("trace_metadata", "prompt_token_count")
    op.drop_index("idx_spans_total_token_count", table_name="spans")
    op.drop_index("idx_spans_total_token_cost", table_name="spans")
    op.drop_column("spans", "total_token_cost")
    op.drop_column("spans", "completion_token_cost")
    op.drop_column("spans", "prompt_token_cost")
    op.drop_column("spans", "total_token_count")
    op.drop_column("spans", "completion_token_count")
    op.drop_column("spans", "prompt_token_count")
    # ### end Alembic commands ###
