"""Sync learning from feedback and outcomes."""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any

from .models import RetrievalEvent, utcnow

if TYPE_CHECKING:
    from .engine import RetrievalLearningEngine
    from ..models import FeedbackRecord, OutcomeRecord


class SyncLearningMixin:
    """Mixin for sync learning functionality."""

    @staticmethod
    def sync_trace(
        engine: "RetrievalLearningEngine",
        trace_id: str | None,
        learning_store: Any,
    ) -> list[RetrievalEvent]:
        """Sync trace feedback to learning tables."""
        if trace_id is None:
            return []

        from .retrieval_events import RetrievalEventsMixin

        events = RetrievalEventsMixin.list_events(engine, trace_id=trace_id)
        if not events:
            return []

        feedback = learning_store.list_feedback(trace_id=trace_id)
        outcomes = learning_store.list_outcomes(trace_id=trace_id)
        synced: list[RetrievalEvent] = []

        for event in events:
            new_feedback = [item for item in feedback if item.id not in event.applied_feedback_ids]
            new_outcomes = [item for item in outcomes if item.id not in event.applied_outcome_ids]

            if not new_feedback and not new_outcomes:
                synced.append(event)
                continue

            delta, accepted_ids, rejected_ids, bad_retrieval = SyncLearningMixin._score_signals(
                engine, new_feedback, new_outcomes
            )
            updated = SyncLearningMixin._apply_signal_delta(
                engine,
                event,
                quality_delta=delta,
                accepted_ids=accepted_ids,
                rejected_ids=rejected_ids,
                bad_retrieval=bad_retrieval,
                feedback_ids=[item.id for item in new_feedback],
                outcome_ids=[item.id for item in new_outcomes],
            )
            SyncLearningMixin._update_learning_tables(
                engine,
                updated,
                quality_delta=delta,
                accepted_ids=accepted_ids,
                rejected_ids=rejected_ids,
            )
            synced.append(updated)

        return synced

    @staticmethod
    def _score_signals(
        engine: "RetrievalLearningEngine",
        feedback: list["FeedbackRecord"],
        outcomes: list["OutcomeRecord"],
    ) -> tuple[float, set[str], set[str], bool]:
        """Score feedback signals."""
        delta = 0.0
        accepted_ids: set[str] = set()
        rejected_ids: set[str] = set()
        bad_retrieval = False

        for item in feedback:
            if item.kind == "thumbs_up":
                delta += 1.0
            elif item.kind == "thumbs_down":
                delta -= 1.0
            elif item.kind == "numeric_rating" and item.rating is not None:
                delta += (float(item.rating) - 3.0) / 2.0
            elif item.kind == "edited_answer":
                delta -= 0.5
            elif item.kind == "suggestion":
                delta += 0.5 if item.suggestion_outcome == "accepted" else -0.5

            if item.corrected_output:
                delta -= 0.25

            failures = set(item.failure_types or [])
            if failures:
                delta -= 0.25 * len(failures)
            if "bad_retrieval" in failures:
                bad_retrieval = True

            metadata = item.metadata or {}
            accepted_ids.update(
                str(value) for value in metadata.get("accepted_contexts", []) if value is not None
            )
            rejected_ids.update(
                str(value) for value in metadata.get("rejected_contexts", []) if value is not None
            )

        for item in outcomes:
            value = item.value
            if item.kind in {"conversion", "resolution"}:
                if value is True:
                    delta += 1.0
                elif value is False:
                    delta -= 1.0
            elif item.kind in {"escalation", "human_takeover"}:
                if value is True:
                    delta -= 1.0
            elif item.kind in {"ctr", "reply_rate"} and isinstance(value, (int, float)):
                delta += max(-1.0, min(1.0, float(value)))

            metadata = item.metadata or {}
            accepted_ids.update(
                str(value) for value in metadata.get("accepted_contexts", []) if value is not None
            )
            rejected_ids.update(
                str(value) for value in metadata.get("rejected_contexts", []) if value is not None
            )

        return delta, accepted_ids, rejected_ids, bad_retrieval

    @staticmethod
    def _apply_signal_delta(
        engine: "RetrievalLearningEngine",
        event: RetrievalEvent,
        *,
        quality_delta: float,
        accepted_ids: set[str],
        rejected_ids: set[str],
        bad_retrieval: bool,
        feedback_ids: list[str],
        outcome_ids: list[str],
    ) -> RetrievalEvent:
        """Apply signal delta to event."""
        updated = RetrievalEvent(
            id=event.id,
            trace_id=event.trace_id,
            session_id=event.session_id,
            workspace_id=event.workspace_id,
            task_type=event.task_type,
            user_segment=event.user_segment,
            prompt_name=event.prompt_name,
            prompt_version=event.prompt_version,
            original_query=event.original_query,
            effective_query=event.effective_query,
            limit=event.limit,
            filter=event.filter,
            chunks=list(event.chunks),
            created_at=event.created_at,
            quality_score=event.quality_score + quality_delta,
            bad_context_detected=event.bad_context_detected or bad_retrieval or bool(rejected_ids),
            synced_at=utcnow(),
            applied_feedback_ids=[*event.applied_feedback_ids, *feedback_ids],
            applied_outcome_ids=[*event.applied_outcome_ids, *outcome_ids],
        )

        with engine._connect() as connection:
            engine._ensure_schema(connection)
            connection.execute(
                f"""
                UPDATE {engine._table("events")}
                SET quality_score = ?, bad_context_detected = ?, synced_at = ?,
                    applied_feedback_ids_json = ?, applied_outcome_ids_json = ?
                WHERE id = ?
                """,
                (
                    updated.quality_score,
                    1 if updated.bad_context_detected else 0,
                    updated.synced_at,
                    json.dumps(updated.applied_feedback_ids, sort_keys=True),
                    json.dumps(updated.applied_outcome_ids, sort_keys=True),
                    updated.id,
                ),
            )
            connection.commit()

        return updated

    @staticmethod
    def _update_learning_tables(
        engine: "RetrievalLearningEngine",
        event: RetrievalEvent,
        *,
        quality_delta: float,
        accepted_ids: set[str],
        rejected_ids: set[str],
    ) -> None:
        """Update learning tables with event feedback."""
        workspace_key = engine._workspace_value(event.workspace_id)
        should_mark_helpful = quality_delta > 0 and not accepted_ids and not rejected_ids
        should_mark_harmful = quality_delta < 0 and not accepted_ids and not rejected_ids

        with engine._connect() as connection:
            engine._ensure_schema(connection)

            for chunk in event.chunks:
                chunk_id = str(chunk["chunk_id"])
                source_key = str(chunk["source_key"])
                document_id = str(chunk["document_id"])
                connector = chunk.get("connector")
                identifiers = {chunk_id, source_key, document_id}

                helpful = should_mark_helpful or bool(accepted_ids & identifiers)
                harmful = should_mark_harmful or bool(rejected_ids & identifiers)
                accepted = bool(accepted_ids & identifiers)
                rejected = bool(rejected_ids & identifiers)

                if not helpful and not harmful and not accepted and not rejected:
                    continue

                SyncLearningMixin._upsert_source_score(
                    engine,
                    connection,
                    workspace_key=workspace_key,
                    source_key=source_key,
                    connector=connector,
                    helpful=1 if helpful else 0,
                    harmful=1 if harmful else 0,
                    accepted=1 if accepted else 0,
                    rejected=1 if rejected else 0,
                )
                SyncLearningMixin._upsert_chunk_score(
                    engine,
                    connection,
                    workspace_key=workspace_key,
                    chunk_id=chunk_id,
                    document_id=document_id,
                    source_key=source_key,
                    helpful=1 if helpful else 0,
                    harmful=1 if harmful else 0,
                )

                if (
                    harmful
                    or rejected
                    or (event.bad_context_detected and not helpful and not accepted)
                ):
                    SyncLearningMixin._upsert_reindex_trigger(
                        engine,
                        connection,
                        workspace_key=workspace_key,
                        source_key=source_key,
                        chunk_size=int(chunk.get("chunk_size_chars") or 0),
                        reason="retrieval_quality_degraded"
                        if event.bad_context_detected
                        else "rejected_context",
                    )

            if event.effective_query != event.original_query:
                from .query_rewrite import QueryRewriteMixin

                QueryRewriteMixin._upsert_query_rewrite(
                    engine,
                    connection,
                    workspace_id=event.workspace_id,
                    task_type=event.task_type,
                    user_segment=event.user_segment,
                    original_query=event.original_query,
                    rewritten_query=event.effective_query,
                    succeeded=quality_delta >= 0,
                )

            SyncLearningMixin._upsert_profile(
                engine,
                connection,
                workspace_key=workspace_key,
                task_key=engine._segment_value(event.task_type),
                segment_key=engine._segment_value(event.user_segment),
                limit_value=event.limit,
                filter_json=engine._filter_json(event.filter),
                succeeded=quality_delta >= 0,
            )

            connection.commit()

    @staticmethod
    def _upsert_source_score(
        engine: "RetrievalLearningEngine",
        connection: Any,
        *,
        workspace_key: str,
        source_key: str,
        connector: str | None,
        helpful: int,
        harmful: int,
        accepted: int,
        rejected: int,
    ) -> None:
        """Upsert source score."""
        row = connection.execute(
            f"SELECT * FROM {engine._table('source_scores')} WHERE workspace_id = ? AND source_key = ?",
            (workspace_key, source_key),
        ).fetchone()
        helpful_count = int(row["helpful_count"]) if row else 0
        harmful_count = int(row["harmful_count"]) if row else 0
        accepted_count = int(row["accepted_count"]) if row else 0
        rejected_count = int(row["rejected_count"]) if row else 0

        helpful_count += helpful
        harmful_count += harmful
        accepted_count += accepted
        rejected_count += rejected
        usefulness_score = (
            helpful_count + (0.5 * accepted_count) - harmful_count - (0.5 * rejected_count)
        )

        payload = (
            workspace_key,
            source_key,
            connector if connector is not None else (row["connector"] if row else None),
            helpful_count,
            harmful_count,
            accepted_count,
            rejected_count,
            usefulness_score,
            utcnow(),
        )

        if engine.dsn:
            connection.execute(
                f"""
                INSERT INTO {engine._table("source_scores")}
                (workspace_id, source_key, connector, helpful_count, harmful_count,
                 accepted_count, rejected_count, usefulness_score, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT (workspace_id, source_key)
                DO UPDATE SET
                    connector = EXCLUDED.connector,
                    helpful_count = EXCLUDED.helpful_count,
                    harmful_count = EXCLUDED.harmful_count,
                    accepted_count = EXCLUDED.accepted_count,
                    rejected_count = EXCLUDED.rejected_count,
                    usefulness_score = EXCLUDED.usefulness_score,
                    last_updated = EXCLUDED.last_updated
                """,
                payload,
            )
        else:
            connection.execute(
                f"""
                INSERT OR REPLACE INTO {engine._table("source_scores")}
                (workspace_id, source_key, connector, helpful_count, harmful_count,
                 accepted_count, rejected_count, usefulness_score, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                payload,
            )

    @staticmethod
    def _upsert_chunk_score(
        engine: "RetrievalLearningEngine",
        connection: Any,
        *,
        workspace_key: str,
        chunk_id: str,
        document_id: str,
        source_key: str,
        helpful: int,
        harmful: int,
    ) -> None:
        """Upsert chunk score."""
        row = connection.execute(
            f"SELECT * FROM {engine._table('chunk_scores')} WHERE workspace_id = ? AND chunk_id = ?",
            (workspace_key, chunk_id),
        ).fetchone()
        helpful_count = int(row["helpful_count"]) if row else 0
        harmful_count = int(row["harmful_count"]) if row else 0

        helpful_count += helpful
        harmful_count += harmful
        usefulness_score = helpful_count - harmful_count

        payload = (
            workspace_key,
            chunk_id,
            document_id,
            source_key,
            helpful_count,
            harmful_count,
            usefulness_score,
            utcnow(),
        )

        if engine.dsn:
            connection.execute(
                f"""
                INSERT INTO {engine._table("chunk_scores")}
                (workspace_id, chunk_id, document_id, source_key, helpful_count, harmful_count, usefulness_score, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT (workspace_id, chunk_id)
                DO UPDATE SET
                    document_id = EXCLUDED.document_id,
                    source_key = EXCLUDED.source_key,
                    helpful_count = EXCLUDED.helpful_count,
                    harmful_count = EXCLUDED.harmful_count,
                    usefulness_score = EXCLUDED.usefulness_score,
                    last_updated = EXCLUDED.last_updated
                """,
                payload,
            )
        else:
            connection.execute(
                f"""
                INSERT OR REPLACE INTO {engine._table("chunk_scores")}
                (workspace_id, chunk_id, document_id, source_key, helpful_count, harmful_count, usefulness_score, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                payload,
            )

    @staticmethod
    def _upsert_profile(
        engine: "RetrievalLearningEngine",
        connection: Any,
        *,
        workspace_key: str,
        task_key: str,
        segment_key: str,
        limit_value: int,
        filter_json: str,
        succeeded: bool,
    ) -> None:
        """Upsert retrieval profile."""
        row = connection.execute(
            f"""
            SELECT success_count, failure_count
            FROM {engine._table("profiles")}
            WHERE workspace_id = ? AND task_type = ? AND user_segment = ?
              AND limit_value = ? AND filter_json = ?
            """,
            (workspace_key, task_key, segment_key, limit_value, filter_json),
        ).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

        payload = (
            workspace_key,
            task_key,
            segment_key,
            limit_value,
            filter_json,
            success_count,
            failure_count,
            utcnow(),
        )

        if engine.dsn:
            connection.execute(
                f"""
                INSERT INTO {engine._table("profiles")}
                (workspace_id, task_type, user_segment, limit_value, filter_json,
                 success_count, failure_count, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT (workspace_id, task_type, user_segment, limit_value, filter_json)
                DO UPDATE SET
                    success_count = EXCLUDED.success_count,
                    failure_count = EXCLUDED.failure_count,
                    last_updated = EXCLUDED.last_updated
                """,
                payload,
            )
        else:
            connection.execute(
                f"""
                INSERT OR REPLACE INTO {engine._table("profiles")}
                (workspace_id, task_type, user_segment, limit_value, filter_json,
                 success_count, failure_count, last_updated)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                payload,
            )

    @staticmethod
    def _upsert_reindex_trigger(
        engine: "RetrievalLearningEngine",
        connection: Any,
        *,
        workspace_key: str,
        source_key: str,
        chunk_size: int,
        reason: str,
    ) -> None:
        """Upsert reindex trigger."""
        from uuid import uuid4

        row = connection.execute(
            f"""
            SELECT * FROM {engine._table("reindex_triggers")}
            WHERE workspace_id = ? AND source_key = ? AND status = 'open'
            ORDER BY updated_at DESC
            LIMIT 1
            """,
            (workspace_key, source_key),
        ).fetchone()

        suggested_chunk_size = max(200, min(1200, int(chunk_size * 0.7) if chunk_size > 0 else 400))
        suggested_overlap = max(40, min(240, int(suggested_chunk_size * 0.2)))
        now = utcnow()

        if row is None:
            connection.execute(
                f"""
                INSERT INTO {engine._table("reindex_triggers")}
                (id, workspace_id, source_key, reason, severity, occurrences, status,
                 suggested_chunk_size, suggested_overlap, metadata_json, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    uuid4().hex,
                    workspace_key,
                    source_key,
                    reason,
                    "high",
                    1,
                    "open",
                    suggested_chunk_size,
                    suggested_overlap,
                    json.dumps({"latest_chunk_size": chunk_size}, sort_keys=True),
                    now,
                    now,
                ),
            )
            return

        connection.execute(
            f"""
            UPDATE {engine._table("reindex_triggers")}
            SET occurrences = ?, severity = ?, reason = ?, suggested_chunk_size = ?,
                suggested_overlap = ?, metadata_json = ?, updated_at = ?
            WHERE id = ?
            """,
            (
                int(row["occurrences"]) + 1,
                "high",
                reason,
                suggested_chunk_size,
                suggested_overlap,
                json.dumps({"latest_chunk_size": chunk_size}, sort_keys=True),
                now,
                row["id"],
            ),
        )
