"""Retrieval learning engine - main orchestration class."""
from __future__ import annotations

from datetime import UTC, datetime, timedelta
import json
import os
from pathlib import Path
import sqlite3
from typing import TYPE_CHECKING, Any

from logos_adk.config import AgentConfig, LearningConfig, RetrievalOpsConfig
from logos_adk.postgres_utils import import_psycopg, resolve_database_url
from logos_adk.rag import Chunk

from .models import (
    QueryRewriteRecord,
    ReindexRequest,
    ReindexTrigger,
    RetrievalEvent,
    RetrievalProfile,
    RetrievalSettings,
    SourceDocumentSnapshot,
    SourceUsefulnessScore,
)

# Import from parent package
from ..models import FeedbackRecord, OutcomeRecord

if TYPE_CHECKING:
    pass


RETRIEVAL_LEARNING_SCHEMA_KIND = "retrieval_learning"
RETRIEVAL_LEARNING_SCHEMA_VERSION = 1


class RetrievalLearningEngine:
    """Persistent retrieval learning for RAG search behavior.
    
    This engine tracks retrieval quality, learns optimal settings per segment,
    manages reindex triggers/requests, and syncs feedback to improve future retrievals.
    """

    def __init__(
        self,
        path: str | None = None,
        table_prefix: str | None = None,
        *,
        dsn: str | None = None,
        source_document_ttl_days: int | None = None,
        max_source_documents_per_source: int | None = None,
        dedupe_window_seconds: int | None = None,
    ) -> None:
        """Initialize retrieval learning engine.
        
        Args:
            path: Path to SQLite database file
            table_prefix: Prefix for database tables
            dsn: PostgreSQL DSN (optional, uses SQLite if not provided)
            source_document_ttl_days: TTL for source document snapshots
            max_source_documents_per_source: Max documents to keep per source
            dedupe_window_seconds: Window for coalescing duplicate requests
        """
        self.path = path or os.getenv("LOGOS_ADK_RETRIEVAL_DB_PATH", ".logos_adk/retrieval_learning.db")
        self.table_prefix = table_prefix or os.getenv("LOGOS_ADK_RETRIEVAL_TABLE_PREFIX", "retrieval_learning")
        self.dsn = resolve_database_url(
            dsn or os.getenv("LOGOS_ADK_RETRIEVAL_DSN"),
            env_vars=("LOGOS_ADK_RETRIEVAL_DSN", "LOGOS_ADK_DATABASE_URL", "DATABASE_URL"),
        )
        self.source_document_ttl_days = max(
            1,
            self._env_int(
                "LOGOS_ADK_RETRIEVAL_SOURCE_TTL_DAYS",
                source_document_ttl_days,
                default=30,
            ),
        )
        self.max_source_documents_per_source = max(
            1,
            self._env_int(
                "LOGOS_ADK_RETRIEVAL_MAX_SOURCE_DOCUMENTS",
                max_source_documents_per_source,
                default=20,
            ),
        )
        self.dedupe_window_seconds = max(
            0,
            self._env_int(
                "LOGOS_ADK_RETRIEVAL_DEDUPE_WINDOW_SECONDS",
                dedupe_window_seconds,
                default=600,
            ),
        )

    @classmethod
    def from_config(
        cls,
        config: AgentConfig | LearningConfig | RetrievalOpsConfig | dict[str, Any] | None,
        *,
        path: str | None = None,
        table_prefix: str | None = None,
    ) -> RetrievalLearningEngine:
        """Build an engine from typed config, raw dict, or env-backed defaults.
        
        Args:
            config: Agent config, Learning config, RetrievalOpsConfig, or dict
            path: Override path for SQLite database
            table_prefix: Override table prefix
            
        Returns:
            Configured RetrievalLearningEngine instance
        """
        retrieval_config: RetrievalOpsConfig | None = None
        if isinstance(config, AgentConfig):
            retrieval_config = config.learning.retrieval
        elif isinstance(config, LearningConfig):
            retrieval_config = config.retrieval
        elif isinstance(config, RetrievalOpsConfig):
            retrieval_config = config
        elif isinstance(config, dict):
            retrieval_data = config.get("retrieval", config)
            if isinstance(retrieval_data, dict):
                retrieval_config = RetrievalOpsConfig(
                    dsn=(
                        None
                        if retrieval_data.get("dsn") in (None, "")
                        else str(retrieval_data.get("dsn"))
                    ),
                    path=str(retrieval_data.get("path", ".logos_adk/retrieval_learning.db")),
                    table_prefix=str(retrieval_data.get("table_prefix", "retrieval_learning")),
                    source_document_ttl_days=int(retrieval_data.get("source_document_ttl_days", 30)),
                    max_source_documents_per_source=int(
                        retrieval_data.get("max_source_documents_per_source", 20)
                    ),
                    dedupe_window_seconds=int(retrieval_data.get("dedupe_window_seconds", 600)),
                )
        if retrieval_config is None:
            return cls(path=path, table_prefix=table_prefix)
        return cls(
            dsn=retrieval_config.dsn,
            path=path or retrieval_config.path,
            table_prefix=table_prefix or retrieval_config.table_prefix,
            source_document_ttl_days=retrieval_config.source_document_ttl_days,
            max_source_documents_per_source=retrieval_config.max_source_documents_per_source,
            dedupe_window_seconds=retrieval_config.dedupe_window_seconds,
        )

    def _env_int(self, name: str, explicit: int | None, *, default: int) -> int:
        """Read integer from env var with explicit override and default."""
        if explicit is not None:
            return explicit
        raw = os.getenv(name)
        if raw is None or not raw.strip():
            return default
        try:
            return int(raw)
        except ValueError:
            return default

    def _table(self, suffix: str) -> str:
        """Generate table name with prefix."""
        return f"{self.table_prefix}_{suffix}"

    def _connect(self) -> sqlite3.Connection | Any:
        """Create database connection (SQLite or PostgreSQL)."""
        if self.dsn:
            psycopg = import_psycopg()
            raw_connection = psycopg.connect(self.dsn, row_factory=psycopg.rows.dict_row)
            return _PostgreSQLCompatConnection(raw_connection)
        db_path = Path(self.path)
        if db_path.parent and not db_path.parent.exists():
            db_path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(self.path)
        connection.row_factory = sqlite3.Row
        return connection

    def _ensure_schema(self, connection: sqlite3.Connection | Any) -> None:
        """Ensure database schema is up to date."""
        # Import schema creation from dedicated module
        from .schema import ensure_schema
        ensure_schema(self, connection)

    def _segment_value(self, value: str | None) -> str:
        """Normalize segment value for storage."""
        return value or "__all__"

    def _workspace_value(self, value: str | None) -> str:
        """Normalize workspace value for storage."""
        return value or "__global__"

    def _normalize_query(self, query: str) -> str:
        """Normalize query for matching (lowercase, single spaces)."""
        return " ".join(query.lower().split())

    def _filter_json(self, filter_value: dict[str, Any] | None) -> str:
        """Serialize filter to canonical JSON string."""
        return json.dumps(filter_value or {}, sort_keys=True)

    def _parse_timestamp(self, value: str | None) -> datetime | None:
        """Parse ISO timestamp string to datetime."""
        if not value:
            return None
        try:
            return datetime.fromisoformat(value)
        except ValueError:
            return None

    def _upsert_values_clause(self) -> str:
        """Get SQL placeholder for upsert values."""
        return "%s" if self.dsn else "?"

    def _is_request_coalescible(self, row: sqlite3.Row) -> bool:
        """Check if reindex request can be coalesced with existing."""
        status = str(row["status"])
        if status == "running":
            return True
        if status != "queued":
            return False
        if self.dedupe_window_seconds <= 0:
            return False
        updated_at = self._parse_timestamp(str(row["updated_at"]))
        if updated_at is None:
            return False
        return updated_at >= datetime.now(UTC) - timedelta(seconds=self.dedupe_window_seconds)

    def _source_key(self, chunk: Chunk) -> str:
        """Extract source key from chunk metadata."""
        metadata = chunk.metadata or {}
        for key in ("source_key", "source", "url", "path", "file", "connector_source", "origin"):
            value = metadata.get(key)
            if value:
                return str(value)
        return chunk.document_id

    def _connector(self, chunk: Chunk) -> str | None:
        """Extract connector/loader from chunk metadata."""
        metadata = chunk.metadata or {}
        for key in ("connector", "loader", "source_type"):
            value = metadata.get(key)
            if value:
                return str(value)
        return None

    def _chunk_payload(self, chunk: Chunk) -> dict[str, Any]:
        """Convert chunk to storable payload."""
        metadata = dict(chunk.metadata or {})
        return {
            "chunk_id": chunk.id,
            "document_id": chunk.document_id,
            "source_key": self._source_key(chunk),
            "connector": self._connector(chunk),
            "metadata": metadata,
            "chunk_size_chars": metadata.get("chunk_size_chars", len(chunk.content)),
        }

    # Delegate methods to mixins - import here to avoid circular imports
    def record_query_rewrite(self, **kwargs: Any) -> QueryRewriteRecord:
        """Record a query rewrite attempt."""
        from .query_rewrite import QueryRewriteMixin
        return QueryRewriteMixin.record_query_rewrite(self, **kwargs)

    def _upsert_query_rewrite(self, connection: Any, **kwargs: Any) -> None:
        """Internal query rewrite upsert."""
        from .query_rewrite import QueryRewriteMixin
        return QueryRewriteMixin._upsert_query_rewrite(self, connection, **kwargs)

    def rewrite_query(self, query: str, **kwargs: Any) -> str:
        """Rewrite query based on learned successful rewrites."""
        from .query_rewrite import QueryRewriteMixin
        return QueryRewriteMixin.rewrite_query(self, query, **kwargs)

    def resolve_settings(self, **kwargs: Any) -> RetrievalSettings:
        """Resolve learned retrieval settings for segment."""
        from .query_rewrite import QueryRewriteMixin
        return QueryRewriteMixin.resolve_settings(self, **kwargs)

    def rerank(self, chunks: list[Chunk], **kwargs: Any) -> list[Chunk]:
        """Rerank chunks based on learned source/chunk scores."""
        from .query_rewrite import QueryRewriteMixin
        return QueryRewriteMixin.rerank(self, chunks, **kwargs)

    def record_retrieval(self, **kwargs: Any) -> RetrievalEvent:
        """Record a retrieval event."""
        from .retrieval_events import RetrievalEventsMixin
        return RetrievalEventsMixin.record_retrieval(self, **kwargs)

    def list_events(self, **kwargs: Any) -> list[RetrievalEvent]:
        """List retrieval events."""
        from .retrieval_events import RetrievalEventsMixin
        return RetrievalEventsMixin.list_events(self, **kwargs)

    def list_query_rewrites(self, **kwargs: Any) -> list[QueryRewriteRecord]:
        """List query rewrite records."""
        from .retrieval_events import RetrievalEventsMixin
        return RetrievalEventsMixin.list_query_rewrites(self, **kwargs)

    def list_source_scores(self, **kwargs: Any) -> list[SourceUsefulnessScore]:
        """List source usefulness scores."""
        from .retrieval_events import RetrievalEventsMixin
        return RetrievalEventsMixin.list_source_scores(self, **kwargs)

    def list_reindex_triggers(self, **kwargs: Any) -> list[ReindexTrigger]:
        """List reindex triggers."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.list_reindex_triggers(self, **kwargs)

    def get_reindex_trigger(self, trigger_id: str) -> ReindexTrigger | None:
        """Get reindex trigger by ID."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.get_reindex_trigger(self, trigger_id)

    def acknowledge_reindex_trigger(self, trigger_id: str, **kwargs: Any) -> ReindexTrigger | None:
        """Acknowledge a reindex trigger."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.acknowledge_reindex_trigger(self, trigger_id, **kwargs)

    def resolve_reindex_trigger(self, trigger_id: str, **kwargs: Any) -> ReindexTrigger | None:
        """Resolve a reindex trigger."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.resolve_reindex_trigger(self, trigger_id, **kwargs)

    def close_reindex_trigger(self, trigger_id: str, **kwargs: Any) -> ReindexTrigger | None:
        """Close a reindex trigger."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.close_reindex_trigger(self, trigger_id, **kwargs)

    def create_force_reindex_request(self, **kwargs: Any) -> ReindexRequest:
        """Create a force reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.create_force_reindex_request(self, **kwargs)

    def find_open_reindex_request(self, **kwargs: Any) -> ReindexRequest | None:
        """Find open reindex request for source."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.find_open_reindex_request(self, **kwargs)

    def list_reindex_requests(self, **kwargs: Any) -> list[ReindexRequest]:
        """List reindex requests."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.list_reindex_requests(self, **kwargs)

    def get_reindex_request(self, request_id: str) -> ReindexRequest | None:
        """Get reindex request by ID."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.get_reindex_request(self, request_id)

    def start_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Start a reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.start_reindex_request(self, request_id, **kwargs)

    def complete_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Complete a reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.complete_reindex_request(self, request_id, **kwargs)

    def fail_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Fail a reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.fail_reindex_request(self, request_id, **kwargs)

    def retry_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Retry a failed reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.retry_reindex_request(self, request_id, **kwargs)

    def resume_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Resume a canceled reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.resume_reindex_request(self, request_id, **kwargs)

    def replay_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Replay a reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.replay_reindex_request(self, request_id, **kwargs)

    def cancel_reindex_request(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Cancel a reindex request."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin.cancel_reindex_request(self, request_id, **kwargs)

    def register_source_document(self, **kwargs: Any) -> SourceDocumentSnapshot:
        """Register a source document snapshot."""
        from .source_documents import SourceDocumentsMixin
        return SourceDocumentsMixin.register_source_document(self, **kwargs)

    def list_source_documents(self, **kwargs: Any) -> list[SourceDocumentSnapshot]:
        """List source document snapshots."""
        from .source_documents import SourceDocumentsMixin
        return SourceDocumentsMixin.list_source_documents(self, **kwargs)

    def prune_source_documents(self, **kwargs: Any) -> int:
        """Prune old source documents."""
        from .source_documents import SourceDocumentsMixin
        return SourceDocumentsMixin.prune_source_documents(self, **kwargs)

    def ops_summary(self, **kwargs: Any) -> dict[str, Any]:
        """Get operations summary."""
        from .profiles import ProfilesMixin
        return ProfilesMixin.ops_summary(self, **kwargs)

    def export_source_scores_report(self, **kwargs: Any) -> str:
        """Export source scores report."""
        from .profiles import ProfilesMixin
        return ProfilesMixin.export_source_scores_report(self, **kwargs)

    def list_profiles(self, **kwargs: Any) -> list[RetrievalProfile]:
        """List retrieval profiles."""
        from .profiles import ProfilesMixin
        return ProfilesMixin.list_profiles(self, **kwargs)

    def sync_trace(self, trace_id: str | None, learning_store: Any) -> list[RetrievalEvent]:
        """Sync trace feedback to learning tables."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin.sync_trace(self, trace_id, learning_store)

    def _score_signals(self, feedback: list[FeedbackRecord], outcomes: list[OutcomeRecord]) -> tuple[float, set[str], set[str], bool]:
        """Score feedback signals."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._score_signals(self, feedback, outcomes)

    def _apply_signal_delta(self, event: RetrievalEvent, **kwargs: Any) -> RetrievalEvent:
        """Apply signal delta to event."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._apply_signal_delta(self, event, **kwargs)

    def _update_learning_tables(self, event: RetrievalEvent, **kwargs: Any) -> None:
        """Update learning tables with event feedback."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._update_learning_tables(self, event, **kwargs)

    def _upsert_source_score(self, connection: Any, **kwargs: Any) -> None:
        """Upsert source score."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._upsert_source_score(self, connection, **kwargs)

    def _upsert_chunk_score(self, connection: Any, **kwargs: Any) -> None:
        """Upsert chunk score."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._upsert_chunk_score(self, connection, **kwargs)

    def _upsert_profile(self, connection: Any, **kwargs: Any) -> None:
        """Upsert retrieval profile."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._upsert_profile(self, connection, **kwargs)

    def _upsert_reindex_trigger(self, connection: Any, **kwargs: Any) -> None:
        """Upsert reindex trigger."""
        from .sync_learning import SyncLearningMixin
        return SyncLearningMixin._upsert_reindex_trigger(self, connection, **kwargs)

    def _update_reindex_trigger_status(self, trigger_id: str, **kwargs: Any) -> ReindexTrigger | None:
        """Update reindex trigger status."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin._update_reindex_trigger_status(self, trigger_id, **kwargs)

    def _update_reindex_request_status(self, request_id: str, **kwargs: Any) -> ReindexRequest | None:
        """Update reindex request status."""
        from .reindex_manager import ReindexManagerMixin
        return ReindexManagerMixin._update_reindex_request_status(self, request_id, **kwargs)

    def _prune_source_documents(self, connection: Any, **kwargs: Any) -> int:
        """Internal prune source documents."""
        from .source_documents import SourceDocumentsMixin
        return SourceDocumentsMixin._prune_source_documents(self, connection, **kwargs)

    def _row_to_event(self, row: sqlite3.Row) -> RetrievalEvent:
        """Convert database row to RetrievalEvent."""
        from .converters import row_to_event
        return row_to_event(row)

    def _row_to_source_score(self, row: sqlite3.Row, workspace_id: str | None) -> SourceUsefulnessScore:
        """Convert database row to SourceUsefulnessScore."""
        from .converters import row_to_source_score
        return row_to_source_score(row, workspace_id)

    def _row_to_reindex_trigger(self, row: sqlite3.Row, workspace_id: str | None) -> ReindexTrigger:
        """Convert database row to ReindexTrigger."""
        from .converters import row_to_reindex_trigger
        return row_to_reindex_trigger(row, workspace_id)

    def _row_to_reindex_request(self, row: sqlite3.Row, workspace_id: str | None) -> ReindexRequest:
        """Convert database row to ReindexRequest."""
        from .converters import row_to_reindex_request
        return row_to_reindex_request(row, workspace_id)

    def _row_to_source_document(self, row: sqlite3.Row, workspace_id: str | None) -> SourceDocumentSnapshot:
        """Convert database row to SourceDocumentSnapshot."""
        from .converters import row_to_source_document
        return row_to_source_document(row, workspace_id)

    def _row_to_profile(self, row: sqlite3.Row) -> RetrievalProfile:
        """Convert database row to RetrievalProfile."""
        from .converters import row_to_profile
        return row_to_profile(row)

    def _source_score_to_dict(self, score: SourceUsefulnessScore) -> dict[str, Any]:
        """Convert SourceUsefulnessScore to dictionary."""
        from .converters import source_score_to_dict
        return source_score_to_dict(score)


class _PostgreSQLCompatConnection:
    """Compatibility wrapper for PostgreSQL connections."""

    def __init__(self, connection: Any) -> None:
        self._connection = connection

    def execute(self, query: str, params: tuple = ()) -> Any:
        result = self._connection.execute(query, params)
        return result

    def commit(self) -> None:
        self._connection.commit()

    def __enter__(self) -> "_PostgreSQLCompatConnection":
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        pass


def _translate_postgresql_query(query: str) -> str:
    """Translate SQLite query to PostgreSQL compatible syntax."""
    # Replace SQLite specific syntax with PostgreSQL equivalents
    query = query.replace("INSERT OR REPLACE", "INSERT")
    query = query.replace("REAL", "DOUBLE PRECISION")
    query = query.replace("INTEGER", "INTEGER")
    return query
