"""Tests for qc_trace.db.writer — batch COPY writer for NormalizedMessage.

These tests require a running PostgreSQL instance. They are skipped
automatically if psycopg is not installed or the database is unreachable.

Run with: pytest tests/test_db_writer.py -v
Requires: scripts/dev-db.sh start
"""

from __future__ import annotations

import asyncio
import json
import os
import uuid
from typing import AsyncIterator

import pytest

# Skip entire module if psycopg is not installed
psycopg = pytest.importorskip("psycopg")

from psycopg import AsyncConnection

from qc_trace.db.migrations import ensure_schema
from qc_trace.db.writer import BatchWriter
from qc_trace.schemas.unified import (
    NormalizedMessage,
    TokenUsage,
    ToolCall,
    ToolResult,
)

DSN = os.environ.get(
    "QC_TRACE_DSN", "postgresql://qc_trace:qc_trace_dev@localhost:5432/qc_trace"
)


def _can_connect() -> bool:
    """Check if the test database is reachable."""
    try:
        with psycopg.connect(DSN, connect_timeout=3) as conn:
            conn.execute("SELECT 1")
        return True
    except Exception:
        return False


requires_db = pytest.mark.skipif(
    not _can_connect(),
    reason="PostgreSQL not available (run scripts/dev-db.sh start)",
)


def _make_message(
    msg_type: str = "assistant",
    session_id: str | None = None,
    content: str = "Hello world",
    tokens: TokenUsage | None = None,
    tool_call: ToolCall | None = None,
    tool_result: ToolResult | None = None,
    model: str = "test-model",
) -> NormalizedMessage:
    return NormalizedMessage(
        id=str(uuid.uuid4()),
        session_id=session_id or str(uuid.uuid4()),
        source="claude_code",
        source_schema_version=1,
        msg_type=msg_type,
        timestamp="2026-02-06T10:00:00Z",
        content=content,
        tokens=tokens or TokenUsage(),
        tool_call=tool_call,
        tool_result=tool_result,
        model=model,
    )


@requires_db
class TestBatchWriter:
    """Integration tests for BatchWriter against a real PostgreSQL instance."""

    @pytest.fixture(autouse=True)
    async def _setup_db(self) -> AsyncIterator[None]:
        """Ensure schema exists and clean up test data after each test."""
        self.conn = await AsyncConnection.connect(DSN)
        await ensure_schema(self.conn)
        yield
        # Clean up all test data
        await self.conn.execute("DELETE FROM tool_results")
        await self.conn.execute("DELETE FROM tool_calls")
        await self.conn.execute("DELETE FROM token_usage")
        await self.conn.execute("DELETE FROM messages")
        await self.conn.execute("DELETE FROM sessions")
        await self.conn.commit()
        await self.conn.close()

    async def test_write_empty_batch(self) -> None:
        writer = BatchWriter(self.conn)
        count = await writer.write([])
        assert count == 0

    async def test_write_single_message(self) -> None:
        msg = _make_message()
        writer = BatchWriter(self.conn)
        count = await writer.write([msg])
        assert count == 1

        result = await self.conn.execute(
            "SELECT id, content, msg_type FROM messages WHERE id = %s",
            (msg.id,),
        )
        row = await result.fetchone()
        assert row is not None
        assert row[0] == msg.id
        assert row[1] == "Hello world"
        assert row[2] == "assistant"

    async def test_write_creates_session(self) -> None:
        session_id = str(uuid.uuid4())
        msg = _make_message(session_id=session_id)
        writer = BatchWriter(self.conn)
        await writer.write([msg])

        result = await self.conn.execute(
            "SELECT id, source, model FROM sessions WHERE id = %s",
            (session_id,),
        )
        row = await result.fetchone()
        assert row is not None
        assert row[0] == session_id
        assert row[1] == "claude_code"
        assert row[2] == "test-model"

    async def test_write_with_token_usage(self) -> None:
        tokens = TokenUsage(input=100, output=50, cached=10, thinking=5)
        msg = _make_message(tokens=tokens)
        writer = BatchWriter(self.conn)
        await writer.write([msg])

        result = await self.conn.execute(
            "SELECT input_tokens, output_tokens, cached_tokens, thinking_tokens "
            "FROM token_usage WHERE message_id = %s",
            (msg.id,),
        )
        row = await result.fetchone()
        assert row is not None
        assert row == (100, 50, 10, 5)

    async def test_write_with_tool_call(self) -> None:
        tc = ToolCall(id="tc-1", name="read_file", input={"path": "/tmp/test"})
        msg = _make_message(msg_type="tool_call", tool_call=tc)
        writer = BatchWriter(self.conn)
        await writer.write([msg])

        result = await self.conn.execute(
            "SELECT tool_id, tool_name, tool_input "
            "FROM tool_calls WHERE message_id = %s",
            (msg.id,),
        )
        row = await result.fetchone()
        assert row is not None
        assert row[0] == "tc-1"
        assert row[1] == "read_file"
        tool_input = row[2] if isinstance(row[2], dict) else json.loads(row[2])
        assert tool_input == {"path": "/tmp/test"}

    async def test_write_with_tool_result(self) -> None:
        tr = ToolResult(call_id="tc-1", output="file contents", status="success")
        msg = _make_message(msg_type="tool_result", tool_result=tr)
        writer = BatchWriter(self.conn)
        await writer.write([msg])

        result = await self.conn.execute(
            "SELECT call_id, output, status "
            "FROM tool_results WHERE message_id = %s",
            (msg.id,),
        )
        row = await result.fetchone()
        assert row is not None
        assert row == ("tc-1", "file contents", "success")

    async def test_write_multiple_messages_same_session(self) -> None:
        session_id = str(uuid.uuid4())
        messages = [
            _make_message(session_id=session_id, msg_type="user", content="hi"),
            _make_message(session_id=session_id, msg_type="assistant", content="hello"),
        ]
        writer = BatchWriter(self.conn)
        count = await writer.write(messages)
        assert count == 2

        result = await self.conn.execute(
            "SELECT COUNT(*) FROM messages WHERE session_id = %s",
            (session_id,),
        )
        row = await result.fetchone()
        assert row[0] == 2

        # Only one session should be created
        result = await self.conn.execute(
            "SELECT COUNT(*) FROM sessions WHERE id = %s",
            (session_id,),
        )
        row = await result.fetchone()
        assert row[0] == 1

    async def test_write_skips_zero_token_usage(self) -> None:
        msg = _make_message(tokens=TokenUsage())
        writer = BatchWriter(self.conn)
        await writer.write([msg])

        result = await self.conn.execute(
            "SELECT COUNT(*) FROM token_usage WHERE message_id = %s",
            (msg.id,),
        )
        row = await result.fetchone()
        assert row[0] == 0

    async def test_session_upsert_updates_last_updated(self) -> None:
        session_id = str(uuid.uuid4())
        msg1 = _make_message(session_id=session_id, model="model-v1")
        msg2 = _make_message(session_id=session_id, model="model-v2")

        writer = BatchWriter(self.conn)
        await writer.write([msg1])
        await writer.write([msg2])

        result = await self.conn.execute(
            "SELECT model FROM sessions WHERE id = %s",
            (session_id,),
        )
        row = await result.fetchone()
        assert row is not None
        # Upsert should update model
        assert row[0] == "model-v2"
