"""SQLite-native trace manager (stdlib sqlite3)."""

from __future__ import annotations

import asyncio
import json
import logging
import re
import sqlite3
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse

from sqlalchemy.engine import make_url

from ..abstractions import (
    EnvironmentEvent,
    LMCAISEvent,
    RuntimeEvent,
    SessionMessageContent,
    SessionTrace,
)
from ..config import CONFIG
from ..storage.base import TraceStorage
from .models import analytics_views

if TYPE_CHECKING:
    from sqlite3 import Connection as SQLiteConnection
else:  # pragma: no cover - runtime strict for typing only
    SQLiteConnection = Any  # type: ignore[assignment]

try:  # pragma: no cover - exercised only when pandas present
    import pandas as pd  # type: ignore
except Exception:  # pragma: no cover
    pd = None  # type: ignore[assignment]

logger = logging.getLogger(__name__)


def _ga(obj: Any, name: str) -> Any:
    """Get attribute from pyclass or dict."""
    if isinstance(obj, dict):
        return obj.get(name)
    return getattr(obj, name, None)


@dataclass(slots=True)
class _ConnectionTarget:
    """Resolved connection target for sqlite."""

    database: str


def _strip_auth_component(url: str) -> tuple[str, str | None]:
    """Remove auth_token query parameter from URL, returning the token separately."""
    parsed = urlparse(url)
    if not parsed.query:
        return url, None

    params = dict(parse_qsl(parsed.query, keep_blank_values=True))
    token = params.pop("auth_token", None)
    query = urlencode(params, doseq=True)
    sanitised = urlunparse(parsed._replace(query=query))
    return sanitised, token


def _resolve_connection_target(db_url: str | None) -> _ConnectionTarget:
    """Normalise the configured database URL (SQLite-only)."""
    url = db_url or CONFIG.db_url
    sanitised, _ = _strip_auth_component(url)

    if sanitised.startswith("sqlite+libsql://"):
        raise RuntimeError("sqlite+libsql scheme is no longer supported; use a file path.")

    # Plain SQLite files: file://, /absolute/path, or relative path
    if sanitised.startswith("file://") or sanitised.startswith("/") or "://" not in sanitised:
        # Strip file:// prefix if present, sqlite3 handles both formats
        db_path = sanitised.replace("file://", "") if sanitised.startswith("file://") else sanitised
        return _ConnectionTarget(database=db_path)

    # Strict to SQLAlchemy URL parsing for anything else we missed.
    try:
        parsed = make_url(sanitised)
        driver = parsed.drivername.lower()
        if driver.startswith("sqlite"):
            database = parsed.database or ""
            if database and database not in {":memory:", ":memory"}:
                # Absolute paths are passed through; relative paths are resolved to cwd
                if database.startswith("/"):
                    db_path = database
                else:
                    db_path = str(Path(database).expanduser().resolve())
            elif database in {":memory:", ":memory"}:
                db_path = ":memory:"
            else:
                raise RuntimeError("SQLite URL missing database path.")
            return _ConnectionTarget(database=db_path)
    except Exception:  # pragma: no cover - defensive guardrail
        logger.debug("Unable to parse db_url via SQLAlchemy", exc_info=True)

    raise RuntimeError(
        "Unsupported tracing database URL. Only local SQLite paths/URLs are supported."
    )


def _json_dumps(value: Any) -> str | None:
    """Serialise Python objects as JSON compatible with the existing schema."""

    def _default(obj: Any):
        if isinstance(obj, datetime):
            return obj.isoformat()
        return str(obj)

    if value is None:
        return None
    return json.dumps(value, separators=(",", ":"), default=_default)


def _maybe_datetime(value: Any) -> Any:
    if value is None or isinstance(value, datetime):
        return value
    if isinstance(value, str):
        try:
            return datetime.fromisoformat(value)
        except ValueError:
            pass
    return value


def _datetime_to_iso(value: Any, *, default_now: bool = False) -> str | None:
    if value is None:
        return datetime.now(UTC).isoformat() if default_now else None
    if isinstance(value, datetime):
        return value.isoformat()
    if isinstance(value, str):
        return value
    return str(value)


def _time_record_values(time_record: Any) -> tuple[float | None, int | None]:
    if isinstance(time_record, dict):
        return time_record.get("event_time"), time_record.get("message_time")
    return getattr(time_record, "event_time", None), getattr(time_record, "message_time", None)


def _message_content_text(content: Any) -> str | None:
    if hasattr(content, "as_text"):
        return content.as_text()
    text = getattr(content, "text", None)
    if text:
        return text
    json_payload = getattr(content, "json_payload", None)
    if json_payload is not None:
        if isinstance(json_payload, str):
            return json_payload
        try:
            return json.dumps(json_payload, default=str)
        except Exception:
            return str(json_payload)
    if hasattr(content, "to_dict"):
        try:
            data = content.to_dict()
            if isinstance(data, dict):
                if data.get("text"):
                    return data["text"]
                if data.get("json_payload") is not None:
                    json_payload = data["json_payload"]
                    if isinstance(json_payload, str):
                        return json_payload
                    return json.dumps(json_payload, default=str)
            return json.dumps(data, default=str)
        except Exception:
            pass
    return str(content)


def _load_json(value: Any) -> Any:
    if value is None or isinstance(value, dict | list):
        return value or {}
    if isinstance(value, str):
        try:
            return json.loads(value)
        except (TypeError, ValueError):
            return {}
    return value


_TABLE_DEFINITIONS: tuple[str, ...] = (
    """
    CREATE TABLE IF NOT EXISTS experiments (
        experiment_id VARCHAR PRIMARY KEY,
        name VARCHAR NOT NULL,
        description TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        configuration TEXT,
        metadata TEXT
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS systems (
        system_id VARCHAR PRIMARY KEY,
        name VARCHAR NOT NULL,
        system_type VARCHAR,
        description TEXT,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        metadata TEXT
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS system_versions (
        version_id VARCHAR PRIMARY KEY,
        system_id VARCHAR NOT NULL,
        version_number VARCHAR NOT NULL,
        commit_hash VARCHAR,
        created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
        configuration TEXT,
        metadata TEXT,
        FOREIGN KEY(system_id) REFERENCES systems(system_id),
        UNIQUE(system_id, version_number)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS experimental_systems (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        experiment_id VARCHAR NOT NULL,
        system_id VARCHAR NOT NULL,
        version_id VARCHAR NOT NULL,
        FOREIGN KEY(experiment_id) REFERENCES experiments(experiment_id),
        FOREIGN KEY(system_id) REFERENCES systems(system_id),
        FOREIGN KEY(version_id) REFERENCES system_versions(version_id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS session_traces (
        session_id VARCHAR PRIMARY KEY,
        created_at DATETIME NOT NULL,
        num_timesteps INTEGER NOT NULL,
        num_events INTEGER NOT NULL,
        num_messages INTEGER NOT NULL,
        metadata TEXT,
        experiment_id VARCHAR,
        embedding VECTOR,
        FOREIGN KEY(experiment_id) REFERENCES experiments(experiment_id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS session_timesteps (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id VARCHAR NOT NULL,
        step_id VARCHAR NOT NULL,
        step_index INTEGER NOT NULL,
        turn_number INTEGER,
        started_at DATETIME,
        completed_at DATETIME,
        num_events INTEGER,
        num_messages INTEGER,
        step_metadata TEXT,
        UNIQUE(session_id, step_id),
        FOREIGN KEY(session_id) REFERENCES session_traces(session_id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS events (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id VARCHAR NOT NULL,
        timestep_id INTEGER,
        event_type VARCHAR NOT NULL,
        system_instance_id VARCHAR,
        event_time FLOAT,
        message_time INTEGER,
        created_at DATETIME,
        model_name VARCHAR,
        provider VARCHAR,
        input_tokens INTEGER,
        output_tokens INTEGER,
        total_tokens INTEGER,
        cost_usd INTEGER,
        latency_ms INTEGER,
        span_id VARCHAR,
        trace_id VARCHAR,
        call_records TEXT,
        reward FLOAT,
        terminated BOOLEAN,
        truncated BOOLEAN,
        system_state_before TEXT,
        system_state_after TEXT,
        metadata TEXT,
        event_metadata TEXT,
        embedding VECTOR,
        CHECK (event_type IN ('cais', 'environment', 'runtime')),
        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
        FOREIGN KEY(timestep_id) REFERENCES session_timesteps(id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS messages (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id VARCHAR NOT NULL,
        timestep_id INTEGER,
        message_type VARCHAR NOT NULL,
        content TEXT NOT NULL,
        timestamp DATETIME,
        event_time FLOAT,
        message_time INTEGER,
        metadata TEXT,
        embedding VECTOR,
        CHECK (message_type IN ('user', 'assistant', 'system', 'tool_use', 'tool_result')),
        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
        FOREIGN KEY(timestep_id) REFERENCES session_timesteps(id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS outcome_rewards (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id VARCHAR NOT NULL,
        objective_key TEXT NOT NULL DEFAULT 'reward',
        total_reward FLOAT NOT NULL,
        achievements_count INTEGER NOT NULL,
        total_steps INTEGER NOT NULL,
        created_at DATETIME NOT NULL,
        reward_metadata TEXT,
        annotation TEXT,
        FOREIGN KEY(session_id) REFERENCES session_traces(session_id)
    )
    """,
    """
    CREATE TABLE IF NOT EXISTS event_rewards (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        event_id INTEGER NOT NULL,
        session_id VARCHAR NOT NULL,
        message_id INTEGER,
        turn_number INTEGER,
        objective_key TEXT NOT NULL DEFAULT 'reward',
        reward_value FLOAT NOT NULL,
        reward_type VARCHAR,
        "key" VARCHAR,
        annotation TEXT,
        source VARCHAR,
        created_at DATETIME NOT NULL,
        FOREIGN KEY(event_id) REFERENCES events(id),
        FOREIGN KEY(session_id) REFERENCES session_traces(session_id),
        FOREIGN KEY(message_id) REFERENCES messages(id)
    )
    """,
)


_INDEX_DEFINITIONS: tuple[str, ...] = (
    "CREATE INDEX IF NOT EXISTS idx_session_created ON session_traces (created_at)",
    "CREATE INDEX IF NOT EXISTS idx_session_experiment ON session_traces (experiment_id)",
    "CREATE INDEX IF NOT EXISTS idx_timestep_session_step ON session_timesteps (session_id, step_id)",
    "CREATE INDEX IF NOT EXISTS idx_timestep_turn ON session_timesteps (turn_number)",
    "CREATE INDEX IF NOT EXISTS idx_event_session_step ON events (session_id, timestep_id)",
    "CREATE INDEX IF NOT EXISTS idx_event_type ON events (event_type)",
    "CREATE INDEX IF NOT EXISTS idx_event_created ON events (created_at)",
    "CREATE INDEX IF NOT EXISTS idx_event_model ON events (model_name)",
    "CREATE INDEX IF NOT EXISTS idx_event_trace ON events (trace_id)",
    "CREATE INDEX IF NOT EXISTS idx_message_session_step ON messages (session_id, timestep_id)",
    "CREATE INDEX IF NOT EXISTS idx_message_type ON messages (message_type)",
    "CREATE INDEX IF NOT EXISTS idx_message_timestamp ON messages (timestamp)",
    "CREATE INDEX IF NOT EXISTS idx_experiment_created ON experiments (created_at)",
    "CREATE INDEX IF NOT EXISTS idx_experiment_name ON experiments (name)",
    "CREATE INDEX IF NOT EXISTS idx_system_name ON systems (name)",
    "CREATE INDEX IF NOT EXISTS idx_system_type ON systems (system_type)",
    "CREATE UNIQUE INDEX IF NOT EXISTS uq_system_version ON system_versions (system_id, version_number)",
    "CREATE INDEX IF NOT EXISTS idx_version_system ON system_versions (system_id)",
    "CREATE INDEX IF NOT EXISTS idx_version_created ON system_versions (created_at)",
    "CREATE UNIQUE INDEX IF NOT EXISTS uq_experiment_system ON experimental_systems (experiment_id, system_id)",
    "CREATE INDEX IF NOT EXISTS idx_experimental_system ON experimental_systems (experiment_id, system_id)",
    "CREATE INDEX IF NOT EXISTS idx_outcome_rewards_session ON outcome_rewards (session_id)",
    "CREATE INDEX IF NOT EXISTS idx_outcome_rewards_session_objective ON outcome_rewards (session_id, objective_key)",
    "CREATE INDEX IF NOT EXISTS idx_outcome_rewards_total ON outcome_rewards (total_reward)",
    "CREATE INDEX IF NOT EXISTS idx_event_rewards_session ON event_rewards (session_id)",
    "CREATE INDEX IF NOT EXISTS idx_event_rewards_event ON event_rewards (event_id)",
    "CREATE INDEX IF NOT EXISTS idx_event_rewards_session_event_objective ON event_rewards (session_id, event_id, objective_key)",
    "CREATE INDEX IF NOT EXISTS idx_event_rewards_type ON event_rewards (reward_type)",
    'CREATE INDEX IF NOT EXISTS idx_event_rewards_key ON event_rewards ("key")',
)


class SQLiteTraceManager(TraceStorage):
    """SQLite-backed trace manager (stdlib sqlite3)."""

    def __init__(
        self,
        db_url: str | None = None,
    ):
        self._target = _resolve_connection_target(db_url)
        self._conn: SQLiteConnection | None = None
        self._conn_lock = asyncio.Lock()
        self._op_lock = asyncio.Lock()
        self._initialized = False

    def _open_connection(self) -> SQLiteConnection:
        """Open a sqlite3 connection for the resolved target."""
        logger.debug("Opening sqlite connection to %s", self._target.database)
        conn = sqlite3.connect(self._target.database, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA synchronous=NORMAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.execute("PRAGMA busy_timeout=5000")
        return conn

    async def initialize(self):
        """Initialise the backend."""
        async with self._conn_lock:
            if self._initialized:
                return

            # Establish a sqlite connection for future operations.
            self._conn = self._open_connection()
            self._ensure_schema()
            self._initialized = True

    async def close(self):
        """Close the sqlite connection."""
        async with self._conn_lock:
            if self._conn:
                logger.debug("Closing sqlite connection to %s", self._target.database)
                self._conn.close()
                self._conn = None
            self._initialized = False

    # ------------------------------------------------------------------
    # Delegated operations (to be swapped with native libsql versions).
    # ------------------------------------------------------------------

    async def insert_session_trace(self, trace: SessionTrace) -> str:
        await self.initialize()

        import logging as _logging

        _logger = _logging.getLogger(__name__)
        _logger.info(
            f"[TRACE_DEBUG] insert_session_trace START: session_id={trace.session_id}, {len(trace.markov_blanket_message_history)} messages"
        )

        session_exists = await self._session_exists(trace.session_id)
        _logger.info(f"[TRACE_DEBUG] Session exists: {session_exists}")

        step_id_map: dict[str, int] = {}

        if session_exists:
            _logger.warning(
                f"[TRACE_DEBUG] Session {trace.session_id} already exists, skipping events/timesteps, only updating messages!"
            )
            # Don't return early - we need to save messages!
            # Just update metadata
            async with self._op_lock:
                conn = self._conn
                assert conn is not None
                conn.execute(
                    "UPDATE session_traces SET metadata = ? WHERE session_id = ?",
                    (_json_dumps(trace.metadata or {}), trace.session_id),
                )
                conn.commit()
            # Skip events and timesteps to ensure idempotency
        else:
            created_at_val = _datetime_to_iso(trace.created_at, default_now=True)

            async with self._op_lock:
                conn = self._conn
                assert conn is not None
                conn.execute(
                    """
                    INSERT INTO session_traces (
                        session_id,
                        created_at,
                        num_timesteps,
                        num_events,
                        num_messages,
                        metadata
                    )
                    VALUES (?, ?, 0, 0, 0, ?)
                    """,
                    (
                        trace.session_id,
                        created_at_val,
                        _json_dumps(trace.metadata or {}),
                    ),
                )
                conn.commit()
                _logger.info("[TRACE_DEBUG] Session row inserted")

            # Only insert timesteps and events if this is a new session
            for step in trace.session_time_steps:
                step_db_id = await self.ensure_timestep(
                    trace.session_id,
                    step_id=_ga(step, "step_id"),
                    step_index=_ga(step, "step_index"),
                    turn_number=_ga(step, "turn_number"),
                    started_at=_ga(step, "timestamp"),
                    completed_at=_ga(step, "completed_at"),
                    metadata=_ga(step, "step_metadata") or {},
                )
                step_id_map[_ga(step, "step_id")] = step_db_id

            for event in trace.event_history:
                step_ref = None
                metadata = _ga(event, "metadata") or {}
                if isinstance(metadata, dict):
                    step_ref = metadata.get("step_id")
                timestep_db_id = step_id_map.get(step_ref) if step_ref else None
                await self.insert_event_row(
                    trace.session_id,
                    timestep_db_id=timestep_db_id,
                    event=event,
                    metadata_override=_ga(event, "metadata") or {},
                )

        import logging as _logging

        _logger = _logging.getLogger(__name__)
        _logger.info(
            f"[TRACE_DEBUG] insert_session_trace: saving {len(trace.markov_blanket_message_history)} messages (session_exists={session_exists})"
        )

        # Only insert messages if this is a new session (for idempotency)
        if not session_exists:
            for idx, msg in enumerate(trace.markov_blanket_message_history):
                metadata = dict(_ga(msg, "metadata") or {})
                step_ref = metadata.get("step_id")
                content_raw = _ga(msg, "content")
                content_value = content_raw
                if isinstance(content_raw, SessionMessageContent):
                    if content_raw.json_payload:
                        metadata.setdefault("json_payload", content_raw.json_payload)
                        content_value = content_raw.json_payload
                    else:
                        content_value = _message_content_text(content_raw)
                        if content_raw.text:
                            metadata.setdefault("text", content_raw.text)
                elif isinstance(content_raw, dict):
                    jp = content_raw.get("json_payload")
                    txt = content_raw.get("text")
                    if jp:
                        metadata.setdefault("json_payload", jp)
                        content_value = jp
                    elif txt:
                        metadata.setdefault("text", txt)
                        content_value = txt
                    else:
                        content_value = json.dumps(content_raw, ensure_ascii=False)
                elif not isinstance(content_value, str):
                    try:
                        content_value = json.dumps(content_value, ensure_ascii=False)
                    except (TypeError, ValueError):
                        content_value = str(content_value)

                msg_type = _ga(msg, "message_type") or ""
                _logger.info(
                    f"[TRACE_DEBUG]   Message {idx + 1}: type={msg_type}, content_len={len(str(content_value))}"
                )

                try:
                    time_rec = _ga(msg, "time_record")
                    if isinstance(time_rec, dict):
                        event_time = time_rec.get("event_time", 0.0)
                        message_time = time_rec.get("message_time")
                    else:
                        event_time, message_time = _time_record_values(time_rec)
                    await self.insert_message_row(
                        trace.session_id,
                        timestep_db_id=step_id_map.get(step_ref) if step_ref else None,
                        message_type=msg_type,
                        content=content_value,
                        event_time=event_time,
                        message_time=message_time,
                        metadata=metadata,
                    )
                    _logger.info(f"[TRACE_DEBUG]   Message {idx + 1}: saved successfully")
                except Exception as exc:
                    _logger.error(
                        f"[TRACE_DEBUG]   Message {idx + 1}: FAILED TO SAVE: {exc}", exc_info=True
                    )
                    raise
        else:
            _logger.info(
                "[TRACE_DEBUG] Skipping message insertion for existing session (idempotency)"
            )

        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            conn.execute(
                "UPDATE session_traces SET num_timesteps = ?, num_events = ?, num_messages = ?, metadata = ? WHERE session_id = ?",
                (
                    len(trace.session_time_steps),
                    len(trace.event_history),
                    len(trace.markov_blanket_message_history),
                    _json_dumps(trace.metadata or {}),
                    trace.session_id,
                ),
            )
            conn.commit()

        return trace.session_id

    async def get_session_trace(self, session_id: str) -> dict[str, Any] | None:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn
            assert conn is not None

            session_cursor = conn.execute(
                """
                SELECT session_id,
                       created_at,
                       num_timesteps,
                       num_events,
                       num_messages,
                       metadata
                FROM session_traces
                WHERE session_id = ?
                """,
                (session_id,),
            )
            session_row = session_cursor.fetchone()
            session_cursor.close()

            if not session_row:
                return None

            session_columns = [
                "session_id",
                "created_at",
                "num_timesteps",
                "num_events",
                "num_messages",
                "metadata",
            ]
            session_data = dict(zip(session_columns, session_row, strict=True))

            timestep_cursor = conn.execute(
                """
                SELECT step_id,
                       step_index,
                       turn_number,
                       started_at,
                       completed_at,
                       step_metadata
                FROM session_timesteps
                WHERE session_id = ?
                ORDER BY step_index ASC
                """,
                (session_id,),
            )
            timestep_rows = timestep_cursor.fetchall()
            timestep_cursor.close()

        return {
            "session_id": session_data["session_id"],
            "created_at": _maybe_datetime(session_data["created_at"]),
            "num_timesteps": session_data["num_timesteps"],
            "num_events": session_data["num_events"],
            "num_messages": session_data["num_messages"],
            "metadata": _load_json(session_data["metadata"]),
            "timesteps": [
                {
                    "step_id": row[0],
                    "step_index": row[1],
                    "turn_number": row[2],
                    "started_at": _maybe_datetime(row[3]),
                    "completed_at": _maybe_datetime(row[4]),
                    "metadata": _load_json(row[5]),
                }
                for row in timestep_rows
            ],
        }

    async def _session_exists(self, session_id: str) -> bool:
        await self.initialize()
        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            cursor = conn.execute(
                "SELECT 1 FROM session_traces WHERE session_id = ?", (session_id,)
            )
            row = cursor.fetchone()
            cursor.close()
            return row is not None

    @staticmethod
    def _normalise_params(params: dict[str, Any] | None) -> dict[str, Any]:
        if not params:
            return {}
        normalised: dict[str, Any] = {}
        for key, value in params.items():
            if isinstance(value, datetime):
                normalised[key] = value.isoformat()
            else:
                normalised[key] = value
        return normalised

    @staticmethod
    def _prepare_query_params(
        query: str, params: dict[str, Any] | list[Any] | tuple[Any, ...]
    ) -> tuple[str, tuple[Any, ...]]:
        if isinstance(params, dict):
            keys: list[str] = []

            def _replace(match: re.Match[str]) -> str:
                key = match.group(1)
                keys.append(key)
                return "?"

            new_query = re.sub(r":([a-zA-Z_][a-zA-Z0-9_]*)", _replace, query)
            if not keys:
                raise ValueError("No named parameters found in query for provided mapping")
            values = tuple(params[key] for key in keys)
            return new_query, values
        if isinstance(params, list | tuple):
            return query, tuple(params)
        raise TypeError("Unsupported parameter type for query execution")

    def _ensure_schema(self) -> None:
        if not self._conn:
            raise RuntimeError("Connection not initialised")

        for ddl in _TABLE_DEFINITIONS:
            self._conn.execute(ddl)
        self._apply_schema_migrations()
        for ddl in _INDEX_DEFINITIONS:
            self._conn.execute(ddl)
        for view_sql in analytics_views.values():
            self._conn.execute(view_sql)
        self._conn.commit()

    def _apply_schema_migrations(self) -> None:
        """Apply forward-compatible schema changes for existing databases."""
        self._migrate_outcome_rewards()
        self._migrate_reward_objective_keys()

    @staticmethod
    def _col_value(row: Any, *, index: int, key: str) -> Any:
        try:
            return row[key]  # type: ignore[index]
        except Exception:
            return row[index]

    def _migrate_outcome_rewards(self) -> None:
        conn = self._conn
        if conn is None:  # pragma: no cover - defensive guard
            raise RuntimeError("Connection not initialised")

        cursor = conn.execute("PRAGMA table_info(outcome_rewards)")
        rows = cursor.fetchall()
        cursor.close()
        if not rows:
            return

        columns: dict[str, str] = {}
        for row in rows:
            name = str(self._col_value(row, index=1, key="name"))
            col_type = self._col_value(row, index=2, key="type")
            columns[name] = (str(col_type) if col_type is not None else "").strip().upper()

        total_reward_type = columns.get("total_reward", "")
        has_annotation = "annotation" in columns
        has_objective_key = "objective_key" in columns

        # If the DB was created with total_reward INTEGER, rebuild the table to
        # update the declared type and add the annotation column in one step.
        if total_reward_type == "INTEGER":
            conn.execute("BEGIN")
            try:
                conn.execute("ALTER TABLE outcome_rewards RENAME TO outcome_rewards_old")
                conn.execute(
                    """
                    CREATE TABLE outcome_rewards (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        session_id VARCHAR NOT NULL,
                        objective_key TEXT NOT NULL DEFAULT 'reward',
                        total_reward FLOAT NOT NULL,
                        achievements_count INTEGER NOT NULL,
                        total_steps INTEGER NOT NULL,
                        created_at DATETIME NOT NULL,
                        reward_metadata TEXT,
                        annotation TEXT,
                        FOREIGN KEY(session_id) REFERENCES session_traces(session_id)
                    )
                    """
                )
                conn.execute(
                    """
                    INSERT INTO outcome_rewards (
                        id,
                        session_id,
                        objective_key,
                        total_reward,
                        achievements_count,
                        total_steps,
                        created_at,
                        reward_metadata,
                        annotation
                    )
                    SELECT
                        id,
                        session_id,
                        'reward',
                        CAST(total_reward AS FLOAT),
                        achievements_count,
                        total_steps,
                        created_at,
                        reward_metadata,
                        NULL
                    FROM outcome_rewards_old
                    """
                )
                conn.execute("DROP TABLE outcome_rewards_old")
                conn.execute("COMMIT")
            except Exception:
                conn.execute("ROLLBACK")
                raise
            return

        # Otherwise, add annotation column if missing.
        if not has_annotation:
            conn.execute("ALTER TABLE outcome_rewards ADD COLUMN annotation TEXT")
        if not has_objective_key:
            conn.execute(
                "ALTER TABLE outcome_rewards ADD COLUMN objective_key TEXT DEFAULT 'reward'"
            )
        conn.execute(
            "UPDATE outcome_rewards SET objective_key = 'reward' WHERE objective_key IS NULL OR objective_key = ''"
        )

    def _migrate_reward_objective_keys(self) -> None:
        conn = self._conn
        if conn is None:  # pragma: no cover - defensive guard
            raise RuntimeError("Connection not initialised")

        cursor = conn.execute("PRAGMA table_info(event_rewards)")
        rows = cursor.fetchall()
        cursor.close()
        if not rows:
            return

        columns = {str(self._col_value(row, index=1, key="name")) for row in rows}
        if "objective_key" not in columns:
            conn.execute("ALTER TABLE event_rewards ADD COLUMN objective_key TEXT DEFAULT 'reward'")
        conn.execute(
            "UPDATE event_rewards SET objective_key = 'reward' WHERE objective_key IS NULL OR objective_key = ''"
        )

    async def query_traces(self, query: str, params: dict[str, Any] | None = None) -> Any:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            normalised = self._normalise_params(params)
            if normalised:
                prepared_query, prepared_params = self._prepare_query_params(query, normalised)
                cursor = conn.execute(prepared_query, prepared_params)
            else:
                cursor = conn.execute(query)
            try:
                description = cursor.description or []
                columns = [col[0] for col in description]
                rows = cursor.fetchall()
            finally:
                cursor.close()

        if not rows:
            if pd is not None:
                return pd.DataFrame(columns=list(columns))  # type: ignore[arg-type]
            return []

        records = [dict(zip(columns, row, strict=True)) for row in rows]
        if pd is not None:
            return pd.DataFrame(records)
        return records

    async def get_model_usage(
        self,
        start_date=None,
        end_date=None,
        model_name=None,
    ) -> Any:
        query = """
            SELECT * FROM model_usage_stats
            WHERE 1=1
        """
        params: dict[str, Any] = {}
        if start_date:
            params["start_date"] = start_date
            query += " AND last_used >= :start_date"
        if end_date:
            params["end_date"] = end_date
            query += " AND first_used <= :end_date"
        if model_name:
            params["model_name"] = model_name
            query += " AND model_name = :model_name"
        query += " ORDER BY usage_count DESC"
        return await self.query_traces(query, params)

    async def delete_session(self, session_id: str) -> bool:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn
            assert conn is not None

            cursor = conn.execute(
                "SELECT 1 FROM session_traces WHERE session_id = ?", (session_id,)
            )
            exists = cursor.fetchone() is not None
            cursor.close()
            if not exists:
                return False

            conn.execute("DELETE FROM event_rewards WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM outcome_rewards WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM messages WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM events WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM session_timesteps WHERE session_id = ?", (session_id,))
            conn.execute("DELETE FROM session_traces WHERE session_id = ?", (session_id,))
            conn.commit()
            return True

    # Experiment helpers -------------------------------------------------
    async def create_experiment(
        self,
        experiment_id: str,
        name: str,
        description: str | None = None,
        configuration: dict[str, Any] | None = None,
    ) -> str:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            conn.execute(
                """
                INSERT INTO experiments (experiment_id, name, description, configuration)
                VALUES (?, ?, ?, ?)
                ON CONFLICT(experiment_id) DO UPDATE SET
                    name = excluded.name,
                    description = excluded.description,
                    configuration = excluded.configuration
                """,
                (
                    experiment_id,
                    name,
                    description,
                    _json_dumps(configuration or {}),
                ),
            )
            conn.commit()
        return experiment_id

    async def link_session_to_experiment(self, session_id: str, experiment_id: str):
        await self.initialize()

        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            conn.execute(
                "UPDATE session_traces SET experiment_id = ? WHERE session_id = ?",
                (experiment_id, session_id),
            )
            conn.commit()

    async def get_sessions_by_experiment(
        self, experiment_id: str, limit: int | None = None
    ) -> list[dict[str, Any]]:
        await self.initialize()

        sql = """
            SELECT session_id,
                   created_at,
                   num_timesteps,
                   num_events,
                   num_messages,
                   metadata
            FROM session_traces
            WHERE experiment_id = ?
            ORDER BY created_at DESC
        """
        params: list[Any] = [experiment_id]
        if limit is not None:
            sql += " LIMIT ?"
            params.append(limit)

        async with self._op_lock:
            conn = self._conn
            assert conn is not None
            cursor = conn.execute(sql, params)
            rows = cursor.fetchall()
            cursor.close()

        return [
            {
                "session_id": row[0],
                "created_at": _maybe_datetime(row[1]),
                "num_timesteps": row[2],
                "num_events": row[3],
                "num_messages": row[4],
                "metadata": _load_json(row[5]),
            }
            for row in rows
        ]

    async def batch_insert_sessions(
        self, traces: list[SessionTrace], batch_size: int | None = None
    ) -> list[str]:
        batch_size = batch_size or CONFIG.batch_size
        inserted: list[str] = []

        for i in range(0, len(traces), batch_size):
            chunk = traces[i : i + batch_size]
            for trace in chunk:
                session_id = await self.insert_session_trace(trace)
                inserted.append(session_id)
        return inserted

    # Incremental helpers -----------------------------------------------
    async def ensure_session(
        self,
        session_id: str,
        *,
        created_at=None,
        metadata=None,
    ) -> None:
        await self.initialize()

        created_at_val = _datetime_to_iso(created_at, default_now=True)
        metadata_json = _json_dumps(metadata or {})

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            conn.execute(
                """
                INSERT INTO session_traces (
                    session_id, created_at, num_timesteps, num_events, num_messages, metadata
                )
                VALUES (?, ?, 0, 0, 0, ?)
                ON CONFLICT(session_id) DO NOTHING
                """,
                (session_id, created_at_val, metadata_json),
            )
            conn.commit()

    async def ensure_timestep(
        self,
        session_id: str,
        *,
        step_id: str,
        step_index: int,
        turn_number: int | None = None,
        started_at=None,
        completed_at=None,
        metadata=None,
    ) -> int:
        await self.initialize()

        started_at_val = _datetime_to_iso(started_at, default_now=True)
        completed_at_val = _datetime_to_iso(completed_at, default_now=False)
        metadata_json = _json_dumps(metadata or {})

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            cur = conn.execute(
                """
                SELECT id FROM session_timesteps
                WHERE session_id = ? AND step_id = ?
                """,
                (session_id, step_id),
            )
            row = cur.fetchone()
            if row:
                return int(row[0])

            cur = conn.execute(
                """
                INSERT INTO session_timesteps (
                    session_id,
                    step_id,
                    step_index,
                    turn_number,
                    started_at,
                    completed_at,
                    num_events,
                    num_messages,
                    step_metadata
                )
                VALUES (?, ?, ?, ?, ?, ?, 0, 0, ?)
                """,
                (
                    session_id,
                    step_id,
                    step_index,
                    turn_number,
                    started_at_val,
                    completed_at_val,
                    metadata_json,
                ),
            )
            timestep_id = int(cur.lastrowid or 0)
            conn.execute(
                """
                UPDATE session_traces
                SET num_timesteps = num_timesteps + 1
                WHERE session_id = ?
                """,
                (session_id,),
            )
            conn.commit()
            return timestep_id

    async def insert_event_row(
        self,
        session_id: str,
        *,
        timestep_db_id: int | None,
        event: Any,
        metadata_override: dict[str, Any] | None = None,
    ) -> int:
        await self.initialize()

        # Handle both pyclass instances and dicts (from __getattr__ on Rust models)
        is_dict = isinstance(event, dict)
        if is_dict:
            event_type = event.get("event_type", "")
        elif not isinstance(event, EnvironmentEvent | LMCAISEvent | RuntimeEvent):
            raise TypeError(f"Unsupported event type for native manager: {type(event)!r}")
        else:
            event_type = ""

        metadata_json = metadata_override or _ga(event, "metadata") or {}
        event_extra_metadata = _ga(event, "event_metadata")
        system_state_before = _ga(event, "system_state_before")
        system_state_after = _ga(event, "system_state_after")

        time_record = _ga(event, "time_record")
        if isinstance(time_record, dict):
            ev_time = time_record.get("event_time", 0.0)
            msg_time = time_record.get("message_time")
        else:
            ev_time, msg_time = _time_record_values(time_record)

        payload: dict[str, Any] = {
            "session_id": session_id,
            "timestep_id": timestep_db_id,
            "system_instance_id": _ga(event, "system_instance_id"),
            "event_time": ev_time,
            "message_time": msg_time,
            "metadata": metadata_json,
            "event_metadata": event_extra_metadata,
            "system_state_before": system_state_before,
            "system_state_after": system_state_after,
        }

        is_cais = isinstance(event, LMCAISEvent) or (is_dict and event_type == "cais")
        is_env = isinstance(event, EnvironmentEvent) or (is_dict and event_type == "environment")
        is_runtime = isinstance(event, RuntimeEvent) or (is_dict and event_type == "runtime")

        if is_cais:
            raw_call_records = _ga(event, "call_records")
            call_records = None
            if raw_call_records:
                call_records = [
                    r
                    if isinstance(r, dict)
                    else (r.to_dict() if hasattr(r, "to_dict") else asdict(r))
                    for r in raw_call_records
                ]
            cost = _ga(event, "cost_usd")
            payload.update(
                {
                    "event_type": "cais",
                    "model_name": _ga(event, "model_name"),
                    "provider": _ga(event, "provider"),
                    "input_tokens": _ga(event, "input_tokens"),
                    "output_tokens": _ga(event, "output_tokens"),
                    "total_tokens": _ga(event, "total_tokens"),
                    "cost_usd": int(cost * 100) if cost is not None else None,
                    "latency_ms": _ga(event, "latency_ms"),
                    "span_id": _ga(event, "span_id"),
                    "trace_id": _ga(event, "trace_id"),
                    "call_records": call_records,
                }
            )
        elif is_env:
            payload.update(
                {
                    "event_type": "environment",
                    "reward": _ga(event, "reward"),
                    "terminated": _ga(event, "terminated"),
                    "truncated": _ga(event, "truncated"),
                }
            )
        elif is_runtime:
            md = _ga(event, "metadata") or {}
            payload.update(
                {
                    "event_type": "runtime",
                    "metadata": {**md, "actions": _ga(event, "actions") or []},
                }
            )

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            cur = conn.execute(
                """
                INSERT INTO events (
                    session_id,
                    timestep_id,
                    event_type,
                    system_instance_id,
                    event_time,
                    message_time,
                    model_name,
                    provider,
                    input_tokens,
                    output_tokens,
                    total_tokens,
                    cost_usd,
                    latency_ms,
                    span_id,
                    trace_id,
                    call_records,
                    reward,
                    terminated,
                    truncated,
                    system_state_before,
                    system_state_after,
                    metadata,
                    event_metadata
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    payload["session_id"],
                    payload["timestep_id"],
                    payload.get("event_type"),
                    payload["system_instance_id"],
                    payload["event_time"],
                    payload["message_time"],
                    payload.get("model_name"),
                    payload.get("provider"),
                    payload.get("input_tokens"),
                    payload.get("output_tokens"),
                    payload.get("total_tokens"),
                    payload.get("cost_usd"),
                    payload.get("latency_ms"),
                    payload.get("span_id"),
                    payload.get("trace_id"),
                    _json_dumps(payload.get("call_records")),
                    payload.get("reward"),
                    payload.get("terminated"),
                    payload.get("truncated"),
                    _json_dumps(payload.get("system_state_before")),
                    _json_dumps(payload.get("system_state_after")),
                    _json_dumps(payload.get("metadata")),
                    _json_dumps(payload.get("event_metadata")),
                ),
            )
            event_id = int(cur.lastrowid or 0)
            conn.execute(
                """
                UPDATE session_traces
                SET num_events = num_events + 1
                WHERE session_id = ?
                """,
                (session_id,),
            )
            if timestep_db_id is not None:
                conn.execute(
                    """
                    UPDATE session_timesteps
                    SET num_events = num_events + 1
                    WHERE id = ?
                    """,
                    (timestep_db_id,),
                )
            conn.commit()
            return event_id

    async def insert_message_row(
        self,
        session_id: str,
        *,
        timestep_db_id: int | None,
        message_type: str,
        content: Any,
        event_time: float | None = None,
        message_time: int | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> int:
        await self.initialize()

        metadata_payload = dict(metadata or {})
        if isinstance(content, SessionMessageContent):
            if content.json_payload:
                metadata_payload.setdefault("json_payload", content.json_payload)
                content_value = content.json_payload
            else:
                content_value = content.as_text()
                if content.text:
                    metadata_payload.setdefault("text", content.text)
        else:
            content_value = content
            if not isinstance(content_value, str):
                try:
                    content_value = json.dumps(content_value, ensure_ascii=False)
                except (TypeError, ValueError):
                    content_value = str(content_value)

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            cur = conn.execute(
                """
                INSERT INTO messages (
                    session_id,
                    timestep_id,
                    message_type,
                    content,
                    event_time,
                    message_time,
                    metadata
                )
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    session_id,
                    timestep_db_id,
                    message_type,
                    content_value,
                    event_time,
                    message_time,
                    _json_dumps(metadata_payload),
                ),
            )
            message_id = int(cur.lastrowid or 0)
            conn.execute(
                """
                UPDATE session_traces
                SET num_messages = num_messages + 1
                WHERE session_id = ?
                """,
                (session_id,),
            )
            if timestep_db_id is not None:
                conn.execute(
                    """
                    UPDATE session_timesteps
                    SET num_messages = num_messages + 1
                    WHERE id = ?
                    """,
                    (timestep_db_id,),
                )
            conn.commit()
            return message_id

    async def insert_outcome_reward(
        self,
        session_id: str,
        *,
        objective_key: str = "reward",
        total_reward: float,
        achievements_count: int,
        total_steps: int,
        reward_metadata: dict | None = None,
        annotation: dict[str, Any] | None = None,
    ) -> int:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            cur = conn.execute(
                """
                INSERT INTO outcome_rewards (
                    session_id,
                    objective_key,
                    total_reward,
                    achievements_count,
                    total_steps,
                    created_at,
                    reward_metadata,
                    annotation
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    session_id,
                    objective_key,
                    total_reward,
                    achievements_count,
                    total_steps,
                    datetime.now(UTC).isoformat(),
                    _json_dumps(reward_metadata),
                    _json_dumps(annotation),
                ),
            )
            conn.commit()
            return int(cur.lastrowid or 0)

    async def insert_event_reward(
        self,
        session_id: str,
        *,
        event_id: int,
        message_id: int | None = None,
        turn_number: int | None = None,
        objective_key: str = "reward",
        reward_value: float = 0.0,
        reward_type: str | None = None,
        key: str | None = None,
        annotation: dict[str, Any] | None = None,
        source: str | None = None,
    ) -> int:
        await self.initialize()

        async with self._op_lock:
            conn = self._conn

            assert conn is not None
            cur = conn.execute(
                """
                INSERT INTO event_rewards (
                    event_id,
                    session_id,
                    message_id,
                    turn_number,
                    objective_key,
                    reward_value,
                    reward_type,
                    key,
                    annotation,
                    source,
                    created_at
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    event_id,
                    session_id,
                    message_id,
                    turn_number,
                    objective_key,
                    reward_value,
                    reward_type,
                    key,
                    _json_dumps(annotation),
                    source,
                    datetime.now(UTC).isoformat(),
                ),
            )
            conn.commit()
            return int(cur.lastrowid or 0)


# Backwards-compat alias
NativeLibsqlTraceManager = SQLiteTraceManager
