"""Centralized online learning coordination for feedback and outcomes."""
from __future__ import annotations

import asyncio
from collections import OrderedDict
from dataclasses import dataclass, field
import threading
from typing import Any, Callable

from .models import FeedbackRecord, OutcomeRecord
from .store import LearningStore


@dataclass
class MemoryLearningTarget:
    """A memory backend + learning engine pair attached to an agent."""

    agent_name: str
    backend: Any
    learning: Any
    scope_session: Any | None = None


@dataclass
class RetrievalSyncTarget:
    """Retrieval learning engine subscribed to trace-level sync."""

    engine: Any
    name: str | None = None


@dataclass
class ModelRoutingSyncTarget:
    """Model routing learning engine subscribed to trace-level sync."""

    engine: Any
    name: str | None = None
    on_synced: Callable[[str, Any | None, Any], None] | None = None


@dataclass
class WorkflowSyncTarget:
    """Workflow learning target reserved for future trace-level sync."""

    engine: Any | None = None
    name: str | None = None
    on_synced: Callable[[str, Any | None, Any | None], None] | None = None


@dataclass
class LearningSubscriberResult:
    """Result of one subscriber invocation."""

    name: str
    category: str
    synced: int = 0
    skipped: bool = False
    skipped_reason: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class LearningSyncResult:
    """Summary of the online sync work triggered for one record/trace."""

    trace_id: str | None
    memory_synced: int = 0
    retrieval_synced: int = 0
    model_routing_synced: int = 0
    workflow_synced: int = 0
    skipped: bool = False
    skipped_reason: str | None = None
    subscriber_results: list[LearningSubscriberResult] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)


class LearningSubscriber:
    """Base subscriber for centralized online learning updates."""

    category: str = "generic"

    def __init__(self, *, name: str | None = None) -> None:
        self.name = name or self.__class__.__name__

    def on_feedback(
        self,
        *,
        record: FeedbackRecord,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        return None

    def on_outcome(
        self,
        *,
        record: OutcomeRecord,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        return self.sync_trace(
            trace_id=record.trace_id,
            run=run,
            store=store,
            coordinator=coordinator,
        )

    def sync_trace(
        self,
        *,
        trace_id: str | None,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        return None


class MemoryLearningSubscriber(LearningSubscriber):
    """Memory sync subscriber triggered by feedback records."""

    category = "memory"

    def __init__(self, target: MemoryLearningTarget) -> None:
        super().__init__(name=target.agent_name)
        self.target = target

    def on_feedback(
        self,
        *,
        record: FeedbackRecord,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        if run is None:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="missing_run")
        with coordinator._lock:
            if record.id in coordinator._processed_feedback_ids:
                return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="duplicate_feedback")
            coordinator._processed_feedback_ids[record.id] = None
            while len(coordinator._processed_feedback_ids) > coordinator._max_processed_ids:
                coordinator._processed_feedback_ids.popitem(last=False)
        signal = coordinator._feedback_signal(record)
        if signal is None:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="no_memory_signal")
        metadata = getattr(run, "metadata", {}) or {}
        entry_ids = metadata.get("memory_entry_ids") or metadata.get("adk_memory_entry_ids") or []
        if not isinstance(entry_ids, list) or not entry_ids:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="no_memory_entries")
        agent_name = getattr(run, "agent_name", None)
        if agent_name is not None and self.target.agent_name != agent_name:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="agent_mismatch")
        session_id = getattr(run, "session_id", None)
        scoped_session_id = session_id
        if callable(self.target.scope_session):
            try:
                scoped_session_id = self.target.scope_session(session_id)
            except Exception:
                scoped_session_id = session_id
        coordinator._run_async_sync(
            self.target.learning.sync_feedback(
                self.target.backend,
                entry_ids=entry_ids,
                session_id=scoped_session_id,
                helpful=signal,
                reason=coordinator._feedback_reason(record),
            )
        )
        return LearningSubscriberResult(name=self.name, category=self.category, synced=1)


class RetrievalLearningSubscriber(LearningSubscriber):
    """Retrieval sync subscriber."""

    category = "retrieval"

    def __init__(self, target: RetrievalSyncTarget) -> None:
        super().__init__(name=target.name or "retrieval")
        self.target = target

    def sync_trace(
        self,
        *,
        trace_id: str | None,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        if trace_id is None:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="missing_trace_id")
        self.target.engine.sync_trace(trace_id, store)
        return LearningSubscriberResult(name=self.name, category=self.category, synced=1)


class ModelRoutingLearningSubscriber(LearningSubscriber):
    """Model routing sync subscriber."""

    category = "model_routing"

    def __init__(self, target: ModelRoutingSyncTarget) -> None:
        super().__init__(name=target.name or "model_routing")
        self.target = target

    def sync_trace(
        self,
        *,
        trace_id: str | None,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        if trace_id is None:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="missing_trace_id")
        self.target.engine.sync_trace(trace_id, store)
        if callable(self.target.on_synced):
            self.target.on_synced(trace_id, run, self.target.engine)
        return LearningSubscriberResult(name=self.name, category=self.category, synced=1)


class WorkflowLearningSubscriber(LearningSubscriber):
    """Workflow sync placeholder for future workflow learning integration."""

    category = "workflow"

    def __init__(self, target: WorkflowSyncTarget) -> None:
        super().__init__(name=target.name or "workflow")
        self.target = target

    def sync_trace(
        self,
        *,
        trace_id: str | None,
        run: Any | None,
        store: LearningStore,
        coordinator: LearningCoordinator,
    ) -> LearningSubscriberResult | None:
        if trace_id is None:
            return LearningSubscriberResult(name=self.name, category=self.category, skipped=True, skipped_reason="missing_trace_id")
        if self.target.engine is not None and hasattr(self.target.engine, "sync_trace"):
            self.target.engine.sync_trace(trace_id, store)
            if callable(self.target.on_synced):
                self.target.on_synced(trace_id, run, self.target.engine)
            return LearningSubscriberResult(name=self.name, category=self.category, synced=1)
        return LearningSubscriberResult(
            name=self.name,
            category=self.category,
            skipped=True,
            skipped_reason="placeholder_no_engine",
            metadata={"placeholder": True},
        )


class LearningCoordinator:
    """Coordinate online updates across learning subsystems."""

    def __init__(
        self,
        store: LearningStore,
        *,
        subscribers: list[LearningSubscriber] | None = None,
        memory_targets: list[MemoryLearningTarget] | None = None,
        retrieval_targets: list[RetrievalSyncTarget] | None = None,
        model_routing_targets: list[ModelRoutingSyncTarget] | None = None,
        workflow_targets: list[WorkflowSyncTarget] | None = None,
    ) -> None:
        self.store = store
        self.subscribers = list(subscribers or [])
        if not self.subscribers:
            self.subscribers.extend(MemoryLearningSubscriber(target) for target in (memory_targets or []))
            self.subscribers.extend(RetrievalLearningSubscriber(target) for target in (retrieval_targets or []))
            self.subscribers.extend(ModelRoutingLearningSubscriber(target) for target in (model_routing_targets or []))
            self.subscribers.extend(WorkflowLearningSubscriber(target) for target in (workflow_targets or []))
        self._processed_feedback_ids: OrderedDict[str, None] = OrderedDict()
        self._max_processed_ids = 50_000
        self._trace_fingerprints: OrderedDict[str, tuple[tuple[str, ...], tuple[str, ...]]] = OrderedDict()
        self._max_trace_fingerprints = 50_000
        self._lock = threading.Lock()

    def after_feedback(self, *, record: FeedbackRecord, run: Any | None) -> LearningSyncResult:
        result = LearningSyncResult(trace_id=record.trace_id)
        for subscriber in self.subscribers:
            try:
                event_result = subscriber.on_feedback(record=record, run=run, store=self.store, coordinator=self)
            except Exception as exc:
                event_result = LearningSubscriberResult(
                    name=subscriber.name,
                    category=subscriber.category,
                    skipped=True,
                    skipped_reason="subscriber_error",
                    metadata={"error": str(exc)},
                )
            if event_result is not None:
                self._merge_subscriber_result(result, event_result)
        trace_result = self.sync_trace(record.trace_id, run=run)
        self._merge_trace_result(result, trace_result)
        return result

    def after_outcome(self, *, record: OutcomeRecord, run: Any | None) -> LearningSyncResult:
        result = LearningSyncResult(trace_id=record.trace_id)
        for subscriber in self.subscribers:
            try:
                event_result = subscriber.on_outcome(record=record, run=run, store=self.store, coordinator=self)
            except Exception as exc:
                event_result = LearningSubscriberResult(
                    name=subscriber.name,
                    category=subscriber.category,
                    skipped=True,
                    skipped_reason="subscriber_error",
                    metadata={"error": str(exc)},
                )
            if event_result is not None and event_result.category == "memory":
                self._merge_subscriber_result(result, event_result)
        trace_result = self.sync_trace(record.trace_id, run=run)
        self._merge_trace_result(result, trace_result)
        return result

    def sync_trace(self, trace_id: str | None, *, run: Any | None = None, force: bool = False) -> LearningSyncResult:
        result = LearningSyncResult(trace_id=trace_id)
        if trace_id is None:
            result.skipped = True
            result.skipped_reason = "missing_trace_id"
            return result
        if run is None:
            run = self.store.get_run_by_trace(trace_id)
        fingerprint = self._trace_fingerprint(trace_id)
        with self._lock:
            previous = self._trace_fingerprints.get(trace_id)
            if not force and previous == fingerprint:
                result.skipped = True
                result.skipped_reason = "idempotent_no_new_signals"
                return result
            self._trace_fingerprints[trace_id] = fingerprint
            while len(self._trace_fingerprints) > self._max_trace_fingerprints:
                self._trace_fingerprints.popitem(last=False)
        for subscriber in self.subscribers:
            try:
                event_result = subscriber.sync_trace(
                    trace_id=trace_id,
                    run=run,
                    store=self.store,
                    coordinator=self,
                )
            except Exception as exc:
                event_result = LearningSubscriberResult(
                    name=subscriber.name,
                    category=subscriber.category,
                    skipped=True,
                    skipped_reason="subscriber_error",
                    metadata={"error": str(exc)},
                )
            if event_result is not None:
                self._merge_subscriber_result(result, event_result)
        if not any(item.synced > 0 for item in result.subscriber_results) and not result.skipped:
            result.metadata["no_active_trace_subscribers"] = True
        result.metadata["feedback_count"] = len(fingerprint[0])
        result.metadata["outcome_count"] = len(fingerprint[1])
        return result

    def _merge_trace_result(self, target: LearningSyncResult, source: LearningSyncResult) -> None:
        target.retrieval_synced += source.retrieval_synced
        target.model_routing_synced += source.model_routing_synced
        target.workflow_synced += source.workflow_synced
        target.subscriber_results.extend(source.subscriber_results)
        target.metadata.update(source.metadata)
        if source.skipped and not target.skipped:
            target.skipped = True
            target.skipped_reason = source.skipped_reason

    def _merge_subscriber_result(self, result: LearningSyncResult, event_result: LearningSubscriberResult) -> None:
        result.subscriber_results.append(event_result)
        if event_result.category == "memory":
            result.memory_synced += event_result.synced
        elif event_result.category == "retrieval":
            result.retrieval_synced += event_result.synced
        elif event_result.category == "model_routing":
            result.model_routing_synced += event_result.synced
        elif event_result.category == "workflow":
            result.workflow_synced += event_result.synced

    def _trace_fingerprint(self, trace_id: str) -> tuple[tuple[str, ...], tuple[str, ...]]:
        feedback_ids = tuple(sorted(item.id for item in self.store.list_feedback(trace_id=trace_id)))
        outcome_ids = tuple(sorted(item.id for item in self.store.list_outcomes(trace_id=trace_id)))
        return feedback_ids, outcome_ids

    def _feedback_signal(self, record: FeedbackRecord) -> bool | None:
        if record.kind == "thumbs_up":
            return True
        if record.kind == "thumbs_down":
            return False
        if record.kind == "edited_answer":
            return False
        if record.kind == "suggestion":
            if record.suggestion_outcome == "accepted":
                return True
            if record.suggestion_outcome == "rejected":
                return False
        if record.kind == "numeric_rating" and record.rating is not None:
            if record.rating >= 4.0:
                return True
            if record.rating <= 2.0:
                return False
        if record.failure_types:
            return False
        if record.corrected_output or record.expected_output:
            return False
        return None

    def _feedback_reason(self, record: FeedbackRecord) -> str:
        suffix = f":{','.join(record.failure_types)}" if record.failure_types else ""
        return f"feedback:{record.kind}{suffix}"

    def _run_async_sync(self, coro: Any) -> Any:
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return asyncio.run(coro)

        result: dict[str, Any] = {}
        error: dict[str, BaseException] = {}

        def _runner() -> None:
            try:
                result["value"] = asyncio.run(coro)
            except BaseException as exc:  # pragma: no cover - defensive bridge
                error["exc"] = exc

        thread = threading.Thread(target=_runner, daemon=True)
        thread.start()
        thread.join()
        if "exc" in error:
            raise error["exc"]
        return result.get("value")


__all__ = [
    "LearningSubscriber",
    "LearningSubscriberResult",
    "LearningCoordinator",
    "LearningSyncResult",
    "MemoryLearningSubscriber",
    "MemoryLearningTarget",
    "ModelRoutingLearningSubscriber",
    "ModelRoutingSyncTarget",
    "RetrievalLearningSubscriber",
    "RetrievalSyncTarget",
    "WorkflowLearningSubscriber",
    "WorkflowSyncTarget",
]
