"""Data access layer for Run and Step CRUD operations.

This module provides async functions for creating, reading, and updating
Run and Step records. It follows the repository pattern where all database
operations go through these functions rather than directly manipulating models.

Design decisions:
- Separate start/end functions match SDK's event model (start event, then end event)
- All functions take session as first argument for dependency injection
- Keyword-only arguments prevent positional argument errors
- Returns None for not-found cases instead of raising exceptions
"""

from __future__ import annotations

from datetime import datetime
from enum import Enum, auto
from typing import Any
from uuid import UUID

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from ..models import Payload, Run, Step


class _Unset(Enum):
    """Sentinel to distinguish 'not provided' from 'explicitly None'."""

    UNSET = auto()


UNSET = _Unset.UNSET


async def create_run(
    session: AsyncSession,
    *,
    id: UUID,
    pipeline_name: str,
    status: str,
    started_at: datetime,
    input_summary: dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
    request_id: str | None = None,
    user_id: str | None = None,
    environment: str | None = None,
) -> Run:
    """Create a new Run record.

    Called when SDK sends a run_start event.

    Args:
        session: Database session
        id: Run UUID (generated by SDK)
        pipeline_name: Name of the pipeline
        status: Initial status (usually "running")
        started_at: When the run started
        input_summary: Summarized input data
        metadata: Additional metadata dict
        request_id: Request ID for tracing (extracted from metadata)
        user_id: User ID for filtering (extracted from metadata)
        environment: Environment name (extracted from metadata)

    Returns:
        The created Run record.
    """
    run = Run(
        id=id,
        pipeline_name=pipeline_name,
        status=status,
        started_at=started_at,
        input_summary=input_summary,
        metadata_=metadata,
        request_id=request_id,
        user_id=user_id,
        environment=environment,
    )
    session.add(run)
    await session.commit()
    # No refresh needed - expire_on_commit=False keeps objects usable
    return run


async def end_run(
    session: AsyncSession,
    *,
    id: UUID,
    status: str,
    ended_at: datetime,
    output_summary: dict[str, Any] | None = None,
    error_message: str | None = None,
) -> Run | None:
    """Update a Run with completion data.

    Called when SDK sends a run_end event.

    Args:
        session: Database session
        id: Run UUID to update
        status: Final status ("success" or "error")
        ended_at: When the run completed
        output_summary: Summarized output data
        error_message: Error message if status is "error"

    Returns:
        The updated Run record, or None if not found.
    """
    run = await session.get(Run, id)
    if run is None:
        return None

    run.status = status
    run.ended_at = ended_at
    run.output_summary = output_summary
    run.error_message = error_message

    await session.commit()
    return run


async def create_step(
    session: AsyncSession,
    *,
    id: UUID,
    run_id: UUID,
    step_name: str,
    step_type: str,
    index: int,
    started_at: datetime,
    status: str = "running",
    input_summary: dict[str, Any] | None = None,
    input_count: int | None = None,
    metadata: dict[str, Any] | None = None,
) -> Step:
    """Create a new Step record.

    Called when SDK sends a step_start event.

    Args:
        session: Database session
        id: Step UUID (generated by SDK)
        run_id: Parent Run UUID
        step_name: Human-readable step name
        step_type: Type of step (filter, rank, llm, etc.)
        index: Step order within run (0-based)
        started_at: When the step started
        status: Initial status (default "running")
        input_summary: Summarized input data
        input_count: Number of input items
        metadata: Additional step metadata

    Returns:
        The created Step record.
    """
    step = Step(
        id=id,
        run_id=run_id,
        step_name=step_name,
        step_type=step_type,
        index=index,
        started_at=started_at,
        status=status,
        input_summary=input_summary,
        input_count=input_count,
        metadata_=metadata,
    )
    session.add(step)
    await session.commit()
    return step


async def end_step(
    session: AsyncSession,
    *,
    id: UUID,
    status: str,
    ended_at: datetime,
    duration_ms: int | None = None,
    output_summary: dict[str, Any] | None = None,
    output_count: int | None = None,
    reasoning: dict[str, Any] | None = None,
    error_message: str | None = None,
) -> Step | None:
    """Update a Step with completion data.

    Called when SDK sends a step_end event.

    Args:
        session: Database session
        id: Step UUID to update
        status: Final status ("success" or "error")
        ended_at: When the step completed
        duration_ms: Execution time in milliseconds
        output_summary: Summarized output data
        output_count: Number of output items
        reasoning: Why decisions were made
        error_message: Error message if status is "error"

    Returns:
        The updated Step record, or None if not found.
    """
    step = await session.get(Step, id)
    if step is None:
        return None

    step.status = status
    step.ended_at = ended_at
    step.duration_ms = duration_ms
    step.output_summary = output_summary
    step.output_count = output_count
    step.reasoning = reasoning
    step.error_message = error_message

    await session.commit()
    return step


async def get_run(
    session: AsyncSession,
    id: UUID,
    *,
    include_steps: bool = True,
) -> Run | None:
    """Get a Run by ID.

    Args:
        session: Database session
        id: Run UUID to fetch
        include_steps: If True, eagerly load associated Steps

    Returns:
        The Run record with steps (if requested), or None if not found.
    """
    if include_steps:
        stmt = select(Run).where(Run.id == id).options(selectinload(Run.steps))
    else:
        stmt = select(Run).where(Run.id == id)

    result = await session.execute(stmt)
    return result.scalar_one_or_none()


async def get_step(session: AsyncSession, id: UUID) -> Step | None:
    """Get a Step by ID.

    Args:
        session: Database session
        id: Step UUID to fetch

    Returns:
        The Step record, or None if not found.
    """
    return await session.get(Step, id)


def _apply_run_filters(
    stmt,
    *,
    pipeline_name: str | None = None,
    status: str | None = None,
    user_id: str | None = None,
    request_id: str | None = None,
    environment: str | None = None,
):
    """Apply run filters to a SQLAlchemy statement.

    Shared by list_runs and count_runs to ensure consistent filtering.
    """
    if pipeline_name is not None:
        stmt = stmt.where(Run.pipeline_name == pipeline_name)
    if status is not None:
        stmt = stmt.where(Run.status == status)
    if user_id is not None:
        stmt = stmt.where(Run.user_id == user_id)
    if request_id is not None:
        stmt = stmt.where(Run.request_id == request_id)
    if environment is not None:
        stmt = stmt.where(Run.environment == environment)
    return stmt


def _apply_step_filters(
    stmt,
    *,
    run_id: UUID | None = None,
    step_type: str | None = None,
    step_name: str | None = None,
    status: str | None = None,
):
    """Apply step filters to a SQLAlchemy statement.

    Shared by list_steps and count_steps to ensure consistent filtering.
    """
    if run_id is not None:
        stmt = stmt.where(Step.run_id == run_id)
    if step_type is not None:
        stmt = stmt.where(Step.step_type == step_type)
    if step_name is not None:
        stmt = stmt.where(Step.step_name == step_name)
    if status is not None:
        stmt = stmt.where(Step.status == status)
    return stmt


async def list_runs(
    session: AsyncSession,
    *,
    pipeline_name: str | None = None,
    status: str | None = None,
    user_id: str | None = None,
    request_id: str | None = None,
    environment: str | None = None,
    limit: int = 100,
    offset: int = 0,
) -> list[Run]:
    """List runs with optional filters.

    Args:
        session: Database session
        pipeline_name: Filter by pipeline name
        status: Filter by status
        user_id: Filter by user ID
        request_id: Filter by request ID
        environment: Filter by environment
        limit: Maximum number of results
        offset: Number of results to skip

    Returns:
        List of Run records (without steps loaded).
    """
    # Secondary sort on id ensures deterministic pagination when timestamps match
    stmt = select(Run).order_by(Run.started_at.desc(), Run.id)
    stmt = _apply_run_filters(
        stmt,
        pipeline_name=pipeline_name,
        status=status,
        user_id=user_id,
        request_id=request_id,
        environment=environment,
    )
    stmt = stmt.limit(limit).offset(offset)

    result = await session.execute(stmt)
    return list(result.scalars().all())


async def list_steps(
    session: AsyncSession,
    *,
    run_id: UUID | None = None,
    step_type: str | None = None,
    step_name: str | None = None,
    status: str | None = None,
    limit: int = 100,
    offset: int = 0,
) -> list[Step]:
    """List steps with optional filters.

    Args:
        session: Database session
        run_id: Filter by run ID
        step_type: Filter by step type
        step_name: Filter by step name
        status: Filter by status
        limit: Maximum number of results
        offset: Number of results to skip

    Returns:
        List of Step records.
    """
    # Secondary sort on id ensures deterministic pagination when timestamps match
    stmt = select(Step).order_by(Step.started_at.desc(), Step.id)
    stmt = _apply_step_filters(
        stmt,
        run_id=run_id,
        step_type=step_type,
        step_name=step_name,
        status=status,
    )
    stmt = stmt.limit(limit).offset(offset)

    result = await session.execute(stmt)
    return list(result.scalars().all())


async def create_payloads(
    session: AsyncSession,
    *,
    run_id: UUID,
    step_id: UUID | None,
    phase: str,
    payloads: dict[str, Any],
) -> list[Payload]:
    """Create Payload records from SDK's _payloads dict.

    The SDK extracts large data (lists >100 items, strings >1KB) into
    a _payloads dict with reference IDs. This function stores each payload
    as a separate record linked to the run/step.

    Args:
        session: Database session
        run_id: Parent Run UUID
        step_id: Parent Step UUID (None for run-level payloads)
        phase: "input" or "output" - which phase the payloads came from
        payloads: Dict of {ref_id: data} from SDK's _payloads field

    Returns:
        List of created Payload records.
    """
    created = []
    for ref_id, data in payloads.items():
        payload = Payload(
            run_id=run_id,
            step_id=step_id,
            ref_id=ref_id,
            phase=phase,
            data=data,
        )
        session.add(payload)
        created.append(payload)

    await session.commit()
    # No refresh needed - expire_on_commit=False keeps objects usable
    # and auto-generated IDs are populated after flush
    return created


async def get_payloads(
    session: AsyncSession,
    *,
    run_id: UUID,
    step_id: UUID | None | _Unset = UNSET,
    phase: str | None = None,
    ref_id: str | None = None,
) -> list[Payload]:
    """Get Payload records for a run/step.

    Args:
        session: Database session
        run_id: Parent Run UUID
        step_id: Filter by step ID. Three behaviors:
            - UNSET (default): Return all payloads (run-level + step-level)
            - None: Return only run-level payloads (step_id IS NULL)
            - UUID: Return only payloads for that specific step
        phase: Filter by phase ("input" or "output")
        ref_id: Filter by specific reference ID

    Returns:
        List of Payload records matching the filters.
    """
    stmt = select(Payload).where(Payload.run_id == run_id)

    if step_id is not UNSET:
        if step_id is None:
            # Explicitly None → filter for run-level payloads (step_id IS NULL)
            stmt = stmt.where(Payload.step_id.is_(None))
        else:
            # UUID provided → filter for that specific step
            stmt = stmt.where(Payload.step_id == step_id)
    # If UNSET, don't filter by step_id → return all payloads

    if phase is not None:
        stmt = stmt.where(Payload.phase == phase)
    if ref_id is not None:
        stmt = stmt.where(Payload.ref_id == ref_id)

    result = await session.execute(stmt)
    return list(result.scalars().all())


async def count_runs(
    session: AsyncSession,
    *,
    pipeline_name: str | None = None,
    status: str | None = None,
    user_id: str | None = None,
    request_id: str | None = None,
    environment: str | None = None,
) -> int:
    """Count runs matching filters (for pagination total).

    Args:
        session: Database session
        pipeline_name: Filter by pipeline name
        status: Filter by status
        user_id: Filter by user ID
        request_id: Filter by request ID
        environment: Filter by environment

    Returns:
        Total count of runs matching the filters.
    """
    from sqlalchemy import func

    stmt = select(func.count()).select_from(Run)
    stmt = _apply_run_filters(
        stmt,
        pipeline_name=pipeline_name,
        status=status,
        user_id=user_id,
        request_id=request_id,
        environment=environment,
    )

    result = await session.execute(stmt)
    return result.scalar_one()


async def count_steps(
    session: AsyncSession,
    *,
    run_id: UUID | None = None,
    step_type: str | None = None,
    step_name: str | None = None,
    status: str | None = None,
) -> int:
    """Count steps matching filters (for pagination total).

    Args:
        session: Database session
        run_id: Filter by run ID
        step_type: Filter by step type
        step_name: Filter by step name
        status: Filter by status

    Returns:
        Total count of steps matching the filters.
    """
    from sqlalchemy import func

    stmt = select(func.count()).select_from(Step)
    stmt = _apply_step_filters(
        stmt,
        run_id=run_id,
        step_type=step_type,
        step_name=step_name,
        status=status,
    )

    result = await session.execute(stmt)
    return result.scalar_one()
