from __future__ import annotations

import asyncio
import json
from datetime import datetime
from typing import Any

import asyncpg

from analysis.safe_btc5.db_store import database_url, split_dsn_ssl


DDL_LPL_SESSIONS = """
CREATE TABLE IF NOT EXISTS lpl_recording_sessions (
    session_id TEXT PRIMARY KEY,
    slug TEXT NOT NULL,
    condition_id TEXT,
    question TEXT,
    outcomes_json JSONB,
    token_ids_json JSONB,
    output_dir TEXT,
    started_at TIMESTAMPTZ NOT NULL,
    ended_at TIMESTAMPTZ,
    status TEXT NOT NULL,
    notes TEXT,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS lpl_recording_sessions_slug_idx
    ON lpl_recording_sessions (slug, started_at DESC);
"""

DDL_LPL_EVENTS = """
CREATE TABLE IF NOT EXISTS lpl_orderbook_events (
    id BIGSERIAL PRIMARY KEY,
    session_id TEXT NOT NULL,
    slug TEXT NOT NULL,
    condition_id TEXT,
    asset_id TEXT,
    outcome TEXT,
    event_type TEXT NOT NULL,
    message_index BIGINT NOT NULL,
    connection_id INTEGER NOT NULL,
    exchange_ts TIMESTAMPTZ,
    received_at_wall TIMESTAMPTZ NOT NULL,
    received_at_monotonic_ns BIGINT NOT NULL,
    raw_json JSONB NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    UNIQUE (session_id, message_index)
);
CREATE INDEX IF NOT EXISTS lpl_orderbook_events_slug_time_idx
    ON lpl_orderbook_events (slug, received_at_wall DESC);
CREATE INDEX IF NOT EXISTS lpl_orderbook_events_session_idx
    ON lpl_orderbook_events (session_id, message_index);
"""

DDL_LPL_SNAPSHOTS = """
CREATE TABLE IF NOT EXISTS lpl_orderbook_snapshots (
    id BIGSERIAL PRIMARY KEY,
    session_id TEXT NOT NULL,
    slug TEXT NOT NULL,
    asset_id TEXT NOT NULL,
    outcome TEXT,
    snapshot_at_wall TIMESTAMPTZ NOT NULL,
    message_index BIGINT NOT NULL,
    best_bid DOUBLE PRECISION,
    best_ask DOUBLE PRECISION,
    bid_levels_json JSONB NOT NULL,
    ask_levels_json JSONB NOT NULL,
    source TEXT NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS lpl_orderbook_snapshots_slug_time_idx
    ON lpl_orderbook_snapshots (slug, snapshot_at_wall DESC);
CREATE INDEX IF NOT EXISTS lpl_orderbook_snapshots_session_idx
    ON lpl_orderbook_snapshots (session_id, message_index);
"""


def dumps(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True)


class LplOrderbookStore:
    def __init__(self, dsn: str | None = None) -> None:
        self.dsn = database_url(dsn)
        self.conn: asyncpg.Connection | None = None
        self._lock = asyncio.Lock()

    async def __aenter__(self) -> "LplOrderbookStore":
        await self.connect()
        await self.init_schema()
        return self

    async def __aexit__(self, *_exc: object) -> None:
        await self.close()

    async def connect(self) -> None:
        clean_dsn, ssl_ctx = split_dsn_ssl(self.dsn)
        self.conn = await asyncpg.connect(clean_dsn, ssl=ssl_ctx)

    async def close(self) -> None:
        if self.conn is not None:
            await self.conn.close()
            self.conn = None

    async def init_schema(self) -> None:
        assert self.conn is not None
        for ddl in (DDL_LPL_SESSIONS, DDL_LPL_EVENTS, DDL_LPL_SNAPSHOTS):
            for stmt in ddl.strip().split(";"):
                stmt = stmt.strip()
                if stmt:
                    await self.conn.execute(stmt)

    async def create_session(
        self,
        *,
        session_id: str,
        slug: str,
        condition_id: str,
        question: str,
        outcomes: list[str],
        token_ids: list[str],
        output_dir: str,
        started_at: datetime,
    ) -> None:
        assert self.conn is not None
        async with self._lock:
            await self.conn.execute(
                """
                INSERT INTO lpl_recording_sessions (
                    session_id, slug, condition_id, question, outcomes_json,
                    token_ids_json, output_dir, started_at, status
                ) VALUES ($1,$2,$3,$4,$5::jsonb,$6::jsonb,$7,$8,'running')
                ON CONFLICT (session_id) DO UPDATE SET
                    status = 'running',
                    updated_at = NOW()
                """,
                session_id,
                slug,
                condition_id,
                question,
                dumps(outcomes),
                dumps(token_ids),
                output_dir,
                started_at,
            )

    async def finish_session(
        self,
        *,
        session_id: str,
        ended_at: datetime,
        status: str,
        notes: str = "",
    ) -> None:
        assert self.conn is not None
        async with self._lock:
            await self.conn.execute(
                """
                UPDATE lpl_recording_sessions
                SET ended_at = $2, status = $3, notes = $4, updated_at = NOW()
                WHERE session_id = $1
                """,
                session_id,
                ended_at,
                status,
                notes,
            )

    async def insert_event(
        self,
        *,
        session_id: str,
        slug: str,
        condition_id: str,
        asset_id: str,
        outcome: str,
        event_type: str,
        message_index: int,
        connection_id: int,
        exchange_ts: datetime | None,
        received_at_wall: datetime,
        received_at_monotonic_ns: int,
        raw_json: dict[str, Any],
    ) -> None:
        assert self.conn is not None
        async with self._lock:
            await self.conn.execute(
                """
                INSERT INTO lpl_orderbook_events (
                    session_id, slug, condition_id, asset_id, outcome, event_type,
                    message_index, connection_id, exchange_ts, received_at_wall,
                    received_at_monotonic_ns, raw_json
                ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12::jsonb)
                ON CONFLICT (session_id, message_index) DO NOTHING
                """,
                session_id,
                slug,
                condition_id,
                asset_id,
                outcome,
                event_type,
                message_index,
                connection_id,
                exchange_ts,
                received_at_wall,
                received_at_monotonic_ns,
                dumps(raw_json),
            )

    async def insert_snapshot(
        self,
        *,
        session_id: str,
        slug: str,
        asset_id: str,
        outcome: str,
        snapshot_at_wall: datetime,
        message_index: int,
        best_bid: float | None,
        best_ask: float | None,
        bid_levels: list[dict[str, float]],
        ask_levels: list[dict[str, float]],
        source: str,
    ) -> None:
        assert self.conn is not None
        async with self._lock:
            await self.conn.execute(
                """
                INSERT INTO lpl_orderbook_snapshots (
                    session_id, slug, asset_id, outcome, snapshot_at_wall,
                    message_index, best_bid, best_ask, bid_levels_json,
                    ask_levels_json, source
                ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9::jsonb,$10::jsonb,$11)
                """,
                session_id,
                slug,
                asset_id,
                outcome,
                snapshot_at_wall,
                message_index,
                best_bid,
                best_ask,
                dumps(bid_levels),
                dumps(ask_levels),
                source,
            )


class NoopLplOrderbookStore:
    """JSONL-only store used when no DATABASE_URL is configured."""

    async def __aenter__(self) -> "NoopLplOrderbookStore":
        return self

    async def __aexit__(self, *_exc: object) -> None:
        return None

    async def create_session(self, **_kwargs: Any) -> None:
        return None

    async def finish_session(self, **_kwargs: Any) -> None:
        return None

    async def insert_event(self, **_kwargs: Any) -> None:
        return None

    async def insert_snapshot(self, **_kwargs: Any) -> None:
        return None