"""Tests for PostgresStore — requires PostgreSQL running via docker compose.

Run with: pytest -m postgres
Skip by default in normal test runs.

Start PostgreSQL:
  cd examples/langchain-postgres && docker compose up -d
"""

from __future__ import annotations

import os

import pytest

from bandito.models import Span, Trace

pytestmark = pytest.mark.postgres

_DEFAULT_URL = "postgresql://bandito:banditodev@localhost:5433/bandito"


def _postgres_available() -> bool:
    try:
        import psycopg
        conn = psycopg.connect(
            os.environ.get("BANDITO_POSTGRES_URL", _DEFAULT_URL),
            autocommit=True,
        )
        conn.close()
        return True
    except Exception:
        return False


@pytest.fixture
def pg_store():
    """Create a PostgresStore with isolated tables for each test."""
    if not _postgres_available():
        pytest.skip("PostgreSQL not available")

    from bandito.storage.postgres import PostgresStore

    url = os.environ.get("BANDITO_POSTGRES_URL", _DEFAULT_URL)
    store = PostgresStore(conninfo=url, auto_create=True)

    yield store

    # Cleanup: truncate tables
    import psycopg
    conn = psycopg.connect(url, autocommit=True)
    with conn.cursor() as cur:
        cur.execute("TRUNCATE bandito_traces")
    conn.close()
    store.close()


class TestWriteAndReadTrace:
    def test_write_and_get(self, pg_store):
        trace = Trace(project="test", trace_id="tr-1", input="hello")
        trace.spans.append(Span(name="llm", kind="llm"))
        pg_store.write_trace(trace)

        result = pg_store.get_trace("tr-1")
        assert result is not None
        assert result["trace_id"] == "tr-1"
        assert result["input"] == "hello"
        assert len(result["spans"]) == 1

    def test_get_missing(self, pg_store):
        assert pg_store.get_trace("nonexistent") is None

    def test_overwrite(self, pg_store):
        trace = Trace(project="test", trace_id="tr-1", input="v1")
        pg_store.write_trace(trace)

        trace2 = Trace(project="test", trace_id="tr-1", input="v2")
        pg_store.write_trace(trace2)

        result = pg_store.get_trace("tr-1")
        assert result["input"] == "v2"

    def test_iter_traces(self, pg_store):
        for i in range(3):
            pg_store.write_trace(Trace(project="test", trace_id=f"tr-{i}"))

        traces = list(pg_store.iter_traces())
        assert len(traces) == 3


class TestDerivedMethods:
    def test_list_projects(self, pg_store):
        pg_store.write_trace(Trace(project="bot-a", trace_id="t-1"))
        pg_store.write_trace(Trace(project="bot-b", trace_id="t-2"))
        pg_store.write_trace(Trace(project="bot-a", trace_id="t-3"))

        assert pg_store.list_projects() == ["bot-a", "bot-b"]

    def test_write_dispatch(self, pg_store):
        trace = Trace(project="test", trace_id="tr-dispatch")
        pg_store.write(trace)
        assert pg_store.get_trace("tr-dispatch") is not None


class TestTableAutoCreation:
    def test_traces_table_exists(self, pg_store):
        """Traces table should be created automatically on store init."""
        import psycopg
        conn = psycopg.connect(
            os.environ.get("BANDITO_POSTGRES_URL", _DEFAULT_URL),
            autocommit=True,
        )
        with conn.cursor() as cur:
            cur.execute("""
                SELECT table_name FROM information_schema.tables
                WHERE table_name = 'bandito_traces'
            """)
            tables = [row[0] for row in cur.fetchall()]
        conn.close()
        assert tables == ["bandito_traces"]
