"""Retrieval events recording and listing."""
from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any
from uuid import uuid4

from logos_adk.rag import Chunk

from .converters import row_to_event
from .models import QueryRewriteRecord, RetrievalEvent, SourceUsefulnessScore

if TYPE_CHECKING:
    from .engine import RetrievalLearningEngine


class RetrievalEventsMixin:
    """Mixin for retrieval events functionality."""

    @staticmethod
    def record_retrieval(
        engine: "RetrievalLearningEngine",
        *,
        trace_id: str | None,
        session_id: str | None,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
        prompt_name: str | None,
        prompt_version: str | None,
        original_query: str,
        effective_query: str,
        limit: int,
        filter: dict[str, Any] | None,
        chunks: list[Chunk],
    ) -> RetrievalEvent:
        """Record a retrieval event."""
        event = RetrievalEvent(
            id=uuid4().hex,
            trace_id=trace_id,
            session_id=session_id,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            prompt_name=prompt_name,
            prompt_version=prompt_version,
            original_query=original_query,
            effective_query=effective_query,
            limit=limit,
            filter=filter,
            chunks=[engine._chunk_payload(chunk) for chunk in chunks],
        )
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            connection.execute(
                f"""
                INSERT INTO {engine._table('events')}
                (id, trace_id, session_id, workspace_id, task_type, user_segment, prompt_name,
                 prompt_version, original_query, effective_query, limit_value, filter_json,
                 chunks_json, quality_score, bad_context_detected, synced_at,
                 applied_feedback_ids_json, applied_outcome_ids_json, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    event.id,
                    event.trace_id,
                    event.session_id,
                    event.workspace_id,
                    event.task_type,
                    event.user_segment,
                    event.prompt_name,
                    event.prompt_version,
                    event.original_query,
                    event.effective_query,
                    event.limit,
                    engine._filter_json(event.filter),
                    json.dumps(event.chunks, sort_keys=True),
                    event.quality_score,
                    1 if event.bad_context_detected else 0,
                    event.synced_at,
                    json.dumps(event.applied_feedback_ids),
                    json.dumps(event.applied_outcome_ids),
                    event.created_at,
                ),
            )
            connection.commit()
        return event

    @staticmethod
    def list_events(
        engine: "RetrievalLearningEngine",
        *,
        trace_id: str | None = None,
    ) -> list[RetrievalEvent]:
        """List retrieval events."""
        params: list[Any] = []
        where_sql = ""
        if trace_id is not None:
            where_sql = "WHERE trace_id = ?"
            params.append(trace_id)
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            rows = connection.execute(
                f"SELECT * FROM {engine._table('events')} {where_sql} ORDER BY created_at DESC",
                params,
            ).fetchall()
        return [row_to_event(row) for row in rows]

    @staticmethod
    def list_query_rewrites(
        engine: "RetrievalLearningEngine",
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> list[QueryRewriteRecord]:
        """List query rewrite records."""
        clauses = ["workspace_id = ?", "task_type = ?", "user_segment = ?"]
        params: list[Any] = [
            engine._workspace_value(workspace_id),
            engine._segment_value(task_type),
            engine._segment_value(user_segment),
        ]
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            rows = connection.execute(
                f"""
                SELECT * FROM {engine._table('query_rewrites')}
                WHERE {' AND '.join(clauses)}
                ORDER BY (success_count - failure_count) DESC, success_count DESC, last_used_at DESC
                """,
                params,
            ).fetchall()
        return [
            QueryRewriteRecord(
                original_query=row["original_query"],
                rewritten_query=row["rewritten_query"],
                workspace_id=workspace_id,
                task_type=task_type,
                user_segment=user_segment,
                success_count=int(row["success_count"]),
                failure_count=int(row["failure_count"]),
                last_used_at=row["last_used_at"],
            )
            for row in rows
        ]

    @staticmethod
    def list_source_scores(
        engine: "RetrievalLearningEngine",
        *,
        workspace_id: str | None = None,
    ) -> list[SourceUsefulnessScore]:
        """List source usefulness scores."""
        workspace_key = engine._workspace_value(workspace_id)
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            rows = connection.execute(
                f"""
                SELECT * FROM {engine._table('source_scores')}
                WHERE workspace_id = ?
                ORDER BY usefulness_score DESC, helpful_count DESC, source_key ASC
                """,
                (workspace_key,),
            ).fetchall()
        from .converters import row_to_source_score
        return [row_to_source_score(row, workspace_id) for row in rows]
