"""Database schema management for retrieval learning."""
from __future__ import annotations

import sqlite3
from typing import Any

from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
)

from .engine import RETRIEVAL_LEARNING_SCHEMA_KIND, RETRIEVAL_LEARNING_SCHEMA_VERSION


def ensure_schema(engine: Any, connection: sqlite3.Connection | Any) -> None:
    """Ensure database schema is up to date.
    
    Args:
        engine: RetrievalLearningEngine instance
        connection: Database connection
    """
    if engine.dsn:
        _ensure_postgresql_schema(engine, connection)
    else:
        _ensure_sqlite_schema(engine, connection)


def _ensure_postgresql_schema(engine: Any, connection: Any) -> None:
    """Create PostgreSQL schema."""
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('query_rewrites')} (
            workspace_id TEXT NOT NULL,
            task_type TEXT NOT NULL,
            user_segment TEXT NOT NULL,
            original_query TEXT NOT NULL,
            rewritten_query TEXT NOT NULL,
            success_count INTEGER NOT NULL DEFAULT 0,
            failure_count INTEGER NOT NULL DEFAULT 0,
            last_used_at TEXT NOT NULL,
            PRIMARY KEY (workspace_id, task_type, user_segment, original_query, rewritten_query)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('events')} (
            id TEXT PRIMARY KEY,
            trace_id TEXT,
            session_id TEXT,
            workspace_id TEXT,
            task_type TEXT,
            user_segment TEXT,
            prompt_name TEXT,
            prompt_version TEXT,
            original_query TEXT NOT NULL,
            effective_query TEXT NOT NULL,
            limit_value INTEGER NOT NULL,
            filter_json TEXT,
            chunks_json TEXT NOT NULL,
            quality_score DOUBLE PRECISION NOT NULL DEFAULT 0,
            bad_context_detected INTEGER NOT NULL DEFAULT 0,
            synced_at TEXT,
            applied_feedback_ids_json TEXT NOT NULL DEFAULT '[]',
            applied_outcome_ids_json TEXT NOT NULL DEFAULT '[]',
            created_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('events')}_trace ON {engine._table('events')} (trace_id, created_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('source_scores')} (
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            connector TEXT,
            helpful_count INTEGER NOT NULL DEFAULT 0,
            harmful_count INTEGER NOT NULL DEFAULT 0,
            accepted_count INTEGER NOT NULL DEFAULT 0,
            rejected_count INTEGER NOT NULL DEFAULT 0,
            usefulness_score DOUBLE PRECISION NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, source_key)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('chunk_scores')} (
            workspace_id TEXT NOT NULL,
            chunk_id TEXT NOT NULL,
            document_id TEXT,
            source_key TEXT NOT NULL,
            helpful_count INTEGER NOT NULL DEFAULT 0,
            harmful_count INTEGER NOT NULL DEFAULT 0,
            usefulness_score DOUBLE PRECISION NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, chunk_id)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('profiles')} (
            workspace_id TEXT NOT NULL,
            task_type TEXT NOT NULL,
            user_segment TEXT NOT NULL,
            limit_value INTEGER NOT NULL,
            filter_json TEXT NOT NULL,
            success_count INTEGER NOT NULL DEFAULT 0,
            failure_count INTEGER NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, task_type, user_segment, limit_value, filter_json)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('reindex_triggers')} (
            id TEXT PRIMARY KEY,
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            reason TEXT NOT NULL,
            severity TEXT NOT NULL,
            occurrences INTEGER NOT NULL DEFAULT 1,
            status TEXT NOT NULL,
            suggested_chunk_size INTEGER,
            suggested_overlap INTEGER,
            metadata_json TEXT NOT NULL DEFAULT '{{}}',
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('triggers')}_source ON {engine._table('reindex_triggers')} (workspace_id, source_key, updated_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('reindex_requests')} (
            id TEXT PRIMARY KEY,
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            reason TEXT NOT NULL,
            requested_by TEXT,
            status TEXT NOT NULL,
            metadata_json TEXT NOT NULL DEFAULT '{{}}',
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('requests')}_source ON {engine._table('reindex_requests')} (workspace_id, source_key, updated_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('source_documents')} (
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            document_id TEXT NOT NULL,
            content TEXT NOT NULL,
            metadata_json TEXT NOT NULL,
            source_ref TEXT,
            updated_at TEXT NOT NULL,
            PRIMARY KEY (workspace_id, source_key, document_id)
        )
        """
    )
    recorded_version = get_postgresql_schema_version(connection, RETRIEVAL_LEARNING_SCHEMA_KIND)
    if recorded_version != RETRIEVAL_LEARNING_SCHEMA_VERSION:
        set_postgresql_schema_version(connection, RETRIEVAL_LEARNING_SCHEMA_KIND, RETRIEVAL_LEARNING_SCHEMA_VERSION)
    connection.commit()


def _ensure_sqlite_schema(engine: Any, connection: sqlite3.Connection) -> None:
    """Create SQLite schema."""
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('query_rewrites')} (
            workspace_id TEXT NOT NULL,
            task_type TEXT NOT NULL,
            user_segment TEXT NOT NULL,
            original_query TEXT NOT NULL,
            rewritten_query TEXT NOT NULL,
            success_count INTEGER NOT NULL DEFAULT 0,
            failure_count INTEGER NOT NULL DEFAULT 0,
            last_used_at TEXT NOT NULL,
            PRIMARY KEY (workspace_id, task_type, user_segment, original_query, rewritten_query)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('events')} (
            id TEXT PRIMARY KEY,
            trace_id TEXT,
            session_id TEXT,
            workspace_id TEXT,
            task_type TEXT,
            user_segment TEXT,
            prompt_name TEXT,
            prompt_version TEXT,
            original_query TEXT NOT NULL,
            effective_query TEXT NOT NULL,
            limit_value INTEGER NOT NULL,
            filter_json TEXT,
            chunks_json TEXT NOT NULL,
            quality_score REAL NOT NULL DEFAULT 0,
            bad_context_detected INTEGER NOT NULL DEFAULT 0,
            synced_at TEXT,
            applied_feedback_ids_json TEXT NOT NULL DEFAULT '[]',
            applied_outcome_ids_json TEXT NOT NULL DEFAULT '[]',
            created_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('events')}_trace ON {engine._table('events')} (trace_id, created_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('source_scores')} (
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            connector TEXT,
            helpful_count INTEGER NOT NULL DEFAULT 0,
            harmful_count INTEGER NOT NULL DEFAULT 0,
            accepted_count INTEGER NOT NULL DEFAULT 0,
            rejected_count INTEGER NOT NULL DEFAULT 0,
            usefulness_score REAL NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, source_key)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('chunk_scores')} (
            workspace_id TEXT NOT NULL,
            chunk_id TEXT NOT NULL,
            document_id TEXT,
            source_key TEXT NOT NULL,
            helpful_count INTEGER NOT NULL DEFAULT 0,
            harmful_count INTEGER NOT NULL DEFAULT 0,
            usefulness_score REAL NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, chunk_id)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('profiles')} (
            workspace_id TEXT NOT NULL,
            task_type TEXT NOT NULL,
            user_segment TEXT NOT NULL,
            limit_value INTEGER NOT NULL,
            filter_json TEXT NOT NULL,
            success_count INTEGER NOT NULL DEFAULT 0,
            failure_count INTEGER NOT NULL DEFAULT 0,
            last_updated TEXT NOT NULL,
            PRIMARY KEY (workspace_id, task_type, user_segment, limit_value, filter_json)
        )
        """
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('reindex_triggers')} (
            id TEXT PRIMARY KEY,
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            reason TEXT NOT NULL,
            severity TEXT NOT NULL,
            occurrences INTEGER NOT NULL DEFAULT 1,
            status TEXT NOT NULL,
            suggested_chunk_size INTEGER,
            suggested_overlap INTEGER,
            metadata_json TEXT NOT NULL DEFAULT '{{}}',
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('triggers')}_source ON {engine._table('reindex_triggers')} (workspace_id, source_key, updated_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('reindex_requests')} (
            id TEXT PRIMARY KEY,
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            reason TEXT NOT NULL,
            requested_by TEXT,
            status TEXT NOT NULL,
            metadata_json TEXT NOT NULL DEFAULT '{{}}',
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
        """
    )
    connection.execute(
        f"CREATE INDEX IF NOT EXISTS idx_{engine._table('requests')}_source ON {engine._table('reindex_requests')} (workspace_id, source_key, updated_at DESC)"
    )
    connection.execute(
        f"""
        CREATE TABLE IF NOT EXISTS {engine._table('source_documents')} (
            workspace_id TEXT NOT NULL,
            source_key TEXT NOT NULL,
            document_id TEXT NOT NULL,
            content TEXT NOT NULL,
            metadata_json TEXT NOT NULL,
            source_ref TEXT,
            updated_at TEXT NOT NULL,
            PRIMARY KEY (workspace_id, source_key, document_id)
        )
        """
    )
    recorded_version = get_sqlite_schema_version(connection, RETRIEVAL_LEARNING_SCHEMA_KIND)
    if recorded_version != RETRIEVAL_LEARNING_SCHEMA_VERSION:
        set_sqlite_schema_version(connection, RETRIEVAL_LEARNING_SCHEMA_KIND, RETRIEVAL_LEARNING_SCHEMA_VERSION)
    connection.commit()
