"""Query rewrite functionality for retrieval learning."""
from __future__ import annotations

import json
import sqlite3
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from logos_adk.rag import Chunk

from .models import QueryRewriteRecord, RetrievalSettings

if TYPE_CHECKING:
    from .engine import RetrievalLearningEngine


def utcnow() -> str:
    """Return current UTC timestamp."""
    return datetime.now(UTC).isoformat()


class QueryRewriteMixin:
    """Mixin for query rewrite functionality."""

    @staticmethod
    def record_query_rewrite(
        engine: "RetrievalLearningEngine",
        *,
        original_query: str,
        rewritten_query: str,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
        succeeded: bool,
    ) -> QueryRewriteRecord:
        """Record a query rewrite attempt.
        
        Args:
            engine: RetrievalLearningEngine instance
            original_query: Original query text
            rewritten_query: Rewritten query text
            workspace_id: Optional workspace identifier
            task_type: Optional task type
            user_segment: Optional user segment
            succeeded: Whether the rewrite was successful
            
        Returns:
            QueryRewriteRecord with updated counts
        """
        workspace_key = engine._workspace_value(workspace_id)
        task_key = engine._segment_value(task_type)
        segment_key = engine._segment_value(user_segment)
        original = engine._normalize_query(original_query)
        rewritten = engine._normalize_query(rewritten_query)
        now = utcnow()
        
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            row = connection.execute(
                f"""
                SELECT success_count, failure_count
                FROM {engine._table('query_rewrites')}
                WHERE workspace_id = ? AND task_type = ? AND user_segment = ?
                  AND original_query = ? AND rewritten_query = ?
                """,
                (workspace_key, task_key, segment_key, original, rewritten),
            ).fetchone()
            success_count = int(row["success_count"]) if row else 0
            failure_count = int(row["failure_count"]) if row else 0
            if succeeded:
                success_count += 1
            else:
                failure_count += 1
            if engine.dsn:
                connection.execute(
                    f"""
                    INSERT INTO {engine._table('query_rewrites')}
                    (workspace_id, task_type, user_segment, original_query, rewritten_query,
                     success_count, failure_count, last_used_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                    ON CONFLICT (workspace_id, task_type, user_segment, original_query, rewritten_query)
                    DO UPDATE SET
                        success_count = EXCLUDED.success_count,
                        failure_count = EXCLUDED.failure_count,
                        last_used_at = EXCLUDED.last_used_at
                    """,
                    (workspace_key, task_key, segment_key, original, rewritten, success_count, failure_count, now),
                )
            else:
                connection.execute(
                    f"""
                    INSERT OR REPLACE INTO {engine._table('query_rewrites')}
                    (workspace_id, task_type, user_segment, original_query, rewritten_query,
                     success_count, failure_count, last_used_at)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    (workspace_key, task_key, segment_key, original, rewritten, success_count, failure_count, now),
                )
            connection.commit()
        
        return QueryRewriteRecord(
            original_query=original,
            rewritten_query=rewritten,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            success_count=success_count,
            failure_count=failure_count,
            last_used_at=now,
        )

    @staticmethod
    def _upsert_query_rewrite(
        engine: "RetrievalLearningEngine",
        connection: sqlite3.Connection | Any,
        *,
        workspace_id: str | None,
        task_type: str | None,
        user_segment: str | None,
        original_query: str,
        rewritten_query: str,
        succeeded: bool,
    ) -> None:
        """Internal method to upsert query rewrite record."""
        workspace_key = engine._workspace_value(workspace_id)
        task_key = engine._segment_value(task_type)
        segment_key = engine._segment_value(user_segment)
        original = engine._normalize_query(original_query)
        rewritten = engine._normalize_query(rewritten_query)
        
        row = connection.execute(
            f"""
            SELECT success_count, failure_count
            FROM {engine._table('query_rewrites')}
            WHERE workspace_id = ? AND task_type = ? AND user_segment = ?
              AND original_query = ? AND rewritten_query = ?
            """,
            (workspace_key, task_key, segment_key, original, rewritten),
        ).fetchone()
        success_count = int(row["success_count"]) if row else 0
        failure_count = int(row["failure_count"]) if row else 0
        if succeeded:
            success_count += 1
        else:
            failure_count += 1
        
        if engine.dsn:
            connection.execute(
                f"""
                INSERT INTO {engine._table('query_rewrites')}
                (workspace_id, task_type, user_segment, original_query, rewritten_query,
                 success_count, failure_count, last_used_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT (workspace_id, task_type, user_segment, original_query, rewritten_query)
                DO UPDATE SET
                    success_count = EXCLUDED.success_count,
                    failure_count = EXCLUDED.failure_count,
                    last_used_at = EXCLUDED.last_used_at
                """,
                (workspace_key, task_key, segment_key, original, rewritten, success_count, failure_count, utcnow()),
            )
        else:
            connection.execute(
                f"""
                INSERT OR REPLACE INTO {engine._table('query_rewrites')}
                (workspace_id, task_type, user_segment, original_query, rewritten_query,
                 success_count, failure_count, last_used_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (workspace_key, task_key, segment_key, original, rewritten, success_count, failure_count, utcnow()),
            )

    @staticmethod
    def rewrite_query(
        engine: "RetrievalLearningEngine",
        query: str,
        *,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> str:
        """Rewrite query based on learned successful rewrites.
        
        Args:
            engine: RetrievalLearningEngine instance
            query: Query to rewrite
            workspace_id: Optional workspace identifier
            task_type: Optional task type
            user_segment: Optional user segment
            
        Returns:
            Rewritten query or original if no successful rewrite found
        """
        normalized = engine._normalize_query(query)
        workspace_key = engine._workspace_value(workspace_id)
        task_candidates = [engine._segment_value(task_type), "__all__"]
        segment_candidates = [engine._segment_value(user_segment), "__all__"]
        workspace_candidates = [workspace_key, "__global__"]
        
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            for workspace_candidate in workspace_candidates:
                for task_candidate in task_candidates:
                    for segment_candidate in segment_candidates:
                        rows = connection.execute(
                            f"""
                            SELECT rewritten_query, success_count, failure_count
                            FROM {engine._table('query_rewrites')}
                            WHERE workspace_id = ? AND task_type = ? AND user_segment = ?
                              AND original_query = ?
                            ORDER BY (success_count - failure_count) DESC, success_count DESC, last_used_at DESC
                            """,
                            (workspace_candidate, task_candidate, segment_candidate, normalized),
                        ).fetchall()
                        if not rows:
                            continue
                        best = rows[0]
                        if int(best["success_count"]) > int(best["failure_count"]):
                            return str(best["rewritten_query"])
        
        return query

    @staticmethod
    def resolve_settings(
        engine: "RetrievalLearningEngine",
        *,
        limit: int,
        filter: dict[str, Any] | None = None,
        workspace_id: str | None = None,
        task_type: str | None = None,
        user_segment: str | None = None,
    ) -> RetrievalSettings:
        """Resolve learned retrieval settings for segment.
        
        Args:
            engine: RetrievalLearningEngine instance
            limit: Default retrieval limit
            filter: Default filter
            workspace_id: Optional workspace identifier
            task_type: Optional task type
            user_segment: Optional user segment
            
        Returns:
            RetrievalSettings with learned values
        """
        workspace_candidates = [engine._workspace_value(workspace_id), "__global__"]
        task_candidates = [engine._segment_value(task_type), "__all__"]
        segment_candidates = [engine._segment_value(user_segment), "__all__"]
        learned_limit = limit
        learned_filter = dict(filter or {}) or None
        
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            for workspace_candidate in workspace_candidates:
                for task_candidate in task_candidates:
                    for segment_candidate in segment_candidates:
                        row = connection.execute(
                            f"""
                            SELECT limit_value, filter_json, success_count, failure_count
                            FROM {engine._table('profiles')}
                            WHERE workspace_id = ? AND task_type = ? AND user_segment = ?
                            ORDER BY (success_count - failure_count) DESC, success_count DESC, last_updated DESC
                            LIMIT 1
                            """,
                            (workspace_candidate, task_candidate, segment_candidate),
                        ).fetchone()
                        if row is None:
                            continue
                        if int(row["success_count"]) <= int(row["failure_count"]):
                            continue
                        learned_limit = max(1, int(row["limit_value"]))
                        if learned_filter is None:
                            raw = json.loads(row["filter_json"]) if row["filter_json"] else {}
                            learned_filter = raw or None
                        return RetrievalSettings(limit=learned_limit, filter=learned_filter)
        
        return RetrievalSettings(limit=learned_limit, filter=learned_filter)

    @staticmethod
    def rerank(
        engine: "RetrievalLearningEngine",
        chunks: list[Chunk],
        *,
        workspace_id: str | None = None,
    ) -> list[Chunk]:
        """Rerank chunks based on learned source/chunk scores.
        
        Args:
            engine: RetrievalLearningEngine instance
            chunks: List of chunks to rerank
            workspace_id: Optional workspace identifier
            
        Returns:
            Reranked list of chunks
        """
        if len(chunks) < 2:
            return chunks
        
        workspace_key = engine._workspace_value(workspace_id)
        ranked: list[tuple[float, Chunk]] = []
        
        with engine._connect() as connection:
            engine._ensure_schema(connection)
            for index, chunk in enumerate(chunks):
                source_key = engine._source_key(chunk)
                source_row = connection.execute(
                    f"SELECT usefulness_score FROM {engine._table('source_scores')} WHERE workspace_id = ? AND source_key = ?",
                    (workspace_key, source_key),
                ).fetchone()
                chunk_row = connection.execute(
                    f"SELECT usefulness_score FROM {engine._table('chunk_scores')} WHERE workspace_id = ? AND chunk_id = ?",
                    (workspace_key, chunk.id),
                ).fetchone()
                score = float(len(chunks) - index) * 0.01
                if source_row is not None:
                    score += float(source_row["usefulness_score"])
                if chunk_row is not None:
                    score += float(chunk_row["usefulness_score"])
                ranked.append((score, chunk))
        
        ranked.sort(key=lambda item: item[0], reverse=True)
        return [chunk for _, chunk in ranked]
