"""
Feedback Collector — Collects all 4 feedback signal types.

PRD §7 — Four signal types:
    7.1 Recall Signal    — on_recall
    7.2 Injection Signal — on_prompt_build
    7.3 Adoption Signal  — on_agent_step
    7.4 Outcome Signal   — on_task_complete

PRD §12 — Integration timing:
    12.1 on_recall          → Recall signal
    12.2 on_prompt_build    → Injection signal
    12.3 on_agent_step      → Adoption signal
    12.4 on_task_complete   → Outcome signal

PRD §18 — Performance: Feedback write latency < 20ms
"""

import time
import logging
from typing import List, Optional, Dict
from openmemo.memory_feedback.feedback_store import (
    FeedbackStore, MemoryFeedbackRecord, WorkflowFeedbackRecord
)
from openmemo.memory_feedback.quality_scorer import QualityScorer

logger = logging.getLogger("openmemo.feedback")


class FeedbackCollector:
    """
    Collects recall / injection / adoption / outcome feedback signals.

    Usage:
        collector = FeedbackCollector(store)

        # §12.1 on recall
        collector.on_recall(memory_ids=["mem_1", "mem_2"],
                            task_id="task_1", agent_id="claw")

        # §12.2 on prompt injection
        collector.on_inject(memory_id="mem_1", task_id="task_1",
                             token_size=120, position="experience_block")

        # §12.3 on agent step (adoption detection)
        collector.on_adopt(memory_id="mem_1", workflow_id="wf_1",
                            task_id="task_1", adopted=True, adoption_rate=0.8)

        # §12.4 on task complete
        collector.on_outcome(task_id="task_1", memory_ids=["mem_1"],
                              workflow_id="wf_1", success=True)
    """

    def __init__(self, store: FeedbackStore, scorer: QualityScorer = None):
        self.store = store
        self.scorer = scorer or QualityScorer(store)

        # In-memory batch: task_id → pending feedback state
        # Allows correlating recall → inject → adopt → outcome for the same task
        self._pending: Dict[str, Dict] = {}

    # ── §7.1 Recall Signal ─────────────────────────────────────────────────────

    def on_recall(self, memory_ids: List[str], task_id: str = "",
                  agent_id: str = "", memory_layer: str = "") -> List[str]:
        """
        PRD §12.1 — Record that memories were recalled.

        Returns list of feedback_ids written.
        """
        t0 = time.time()
        ids = []
        for mid in memory_ids:
            fb = MemoryFeedbackRecord(
                memory_id=mid,
                task_id=task_id,
                agent_id=agent_id,
                memory_layer=memory_layer,
                recalled=True,
            )
            self.store.record_memory_feedback(fb)
            self.scorer.update_memory_profile(mid, recalled=True)
            ids.append(fb.feedback_id)

            # Track pending state
            key = f"{task_id}:{mid}"
            self._pending[key] = {
                "memory_id": mid,
                "task_id": task_id,
                "agent_id": agent_id,
                "recalled": True,
                "injected": False,
                "adopted": False,
            }

        elapsed = (time.time() - t0) * 1000
        logger.debug("on_recall: %d memories task=%s latency=%.1fms",
                     len(memory_ids), task_id, elapsed)
        return ids

    # ── §7.2 Injection Signal ──────────────────────────────────────────────────

    def on_inject(self, memory_id: str, task_id: str = "",
                  token_size: int = 0, position: str = "",
                  agent_id: str = "") -> str:
        """
        PRD §12.2 — Record that a memory was injected into a prompt.
        """
        t0 = time.time()
        fb = MemoryFeedbackRecord(
            memory_id=memory_id,
            task_id=task_id,
            agent_id=agent_id,
            injected=True,
            token_size=token_size,
            position=position,
        )
        self.store.record_memory_feedback(fb)
        self.scorer.update_memory_profile(memory_id, injected=True)

        key = f"{task_id}:{memory_id}"
        if key in self._pending:
            self._pending[key]["injected"] = True
        else:
            self._pending[key] = {"memory_id": memory_id, "task_id": task_id,
                                   "injected": True, "recalled": False, "adopted": False}

        elapsed = (time.time() - t0) * 1000
        logger.debug("on_inject: mem=%s task=%s tokens=%d latency=%.1fms",
                     memory_id, task_id, token_size, elapsed)
        return fb.feedback_id

    def on_inject_block(self, task_id: str, block_map: Dict[str, int],
                         position_prefix: str = "", agent_id: str = "") -> List[str]:
        """
        Bulk inject: block_map = {memory_id: token_size, ...}
        """
        ids = []
        for mid, tokens in block_map.items():
            fid = self.on_inject(mid, task_id=task_id, token_size=tokens,
                                  position=position_prefix, agent_id=agent_id)
            ids.append(fid)
        return ids

    # ── §7.3 Adoption Signal ───────────────────────────────────────────────────

    def on_adopt(self, task_id: str, memory_id: str = "",
                  workflow_id: str = "", adopted: bool = False,
                  adoption_rate: float = 0.0, agent_id: str = "") -> str:
        """
        PRD §12.3 — Record adoption signal from agent step observation.

        PRD §7.3 adoption detection:
            If agent's execution steps match recommended workflow steps → adopted=True
        """
        t0 = time.time()

        if memory_id:
            fb = MemoryFeedbackRecord(
                memory_id=memory_id,
                task_id=task_id,
                agent_id=agent_id,
                adopted=adopted,
            )
            self.store.record_memory_feedback(fb)
            self.scorer.update_memory_profile(memory_id, adopted=adopted)

            key = f"{task_id}:{memory_id}"
            if key in self._pending:
                self._pending[key]["adopted"] = adopted

        fid = ""
        if workflow_id:
            wfb = WorkflowFeedbackRecord(
                workflow_id=workflow_id,
                task_id=task_id,
                agent_id=agent_id,
                recommended=True,
                adopted=adopted,
                adoption_rate=adoption_rate,
            )
            self.store.record_workflow_feedback(wfb)
            fid = wfb.feedback_id

        elapsed = (time.time() - t0) * 1000
        logger.debug("on_adopt: mem=%s wf=%s adopted=%s latency=%.1fms",
                     memory_id, workflow_id, adopted, elapsed)
        return fid

    # ── §7.4 Outcome Signal ────────────────────────────────────────────────────

    def on_outcome(self, task_id: str, success: bool = False,
                    memory_ids: List[str] = None, workflow_id: str = "",
                    execution_time: float = 0.0, agent_id: str = "",
                    reuse_detected: bool = False) -> List[str]:
        """
        PRD §12.4 — Record task outcome (success/failure) for all related memories.
        """
        t0 = time.time()
        ids = []
        memory_ids = memory_ids or []

        for mid in memory_ids:
            fb = MemoryFeedbackRecord(
                memory_id=mid,
                task_id=task_id,
                agent_id=agent_id,
                success=success,
            )
            self.store.record_memory_feedback(fb)
            self.scorer.update_memory_profile(mid, success=success)
            ids.append(fb.feedback_id)

        if workflow_id:
            wfb = WorkflowFeedbackRecord(
                workflow_id=workflow_id,
                task_id=task_id,
                agent_id=agent_id,
                success=success,
                completion_time=execution_time,
            )
            self.store.record_workflow_feedback(wfb)
            ids.append(wfb.feedback_id)

        # Clean up pending state
        for mid in memory_ids:
            key = f"{task_id}:{mid}"
            self._pending.pop(key, None)

        elapsed = (time.time() - t0) * 1000
        logger.debug("on_outcome: task=%s success=%s mems=%d latency=%.1fms",
                     task_id, success, len(memory_ids), elapsed)
        return ids

    # ── §7.3 Adoption Detection Helper ────────────────────────────────────────

    def detect_adoption(self, step_text: str,
                         workflow_keywords: List[str]) -> bool:
        """
        PRD §7.3 — Detect if an agent step text indicates workflow adoption.

        Matches keywords from recommended workflow steps in the agent's output.
        """
        if not workflow_keywords or not step_text:
            return False
        tl = step_text.lower()
        matches = sum(1 for kw in workflow_keywords if kw.lower() in tl)
        adoption_rate = matches / len(workflow_keywords)
        return adoption_rate >= 0.3  # 30% keyword match threshold

    def detect_adoption_rate(self, step_text: str,
                               workflow_keywords: List[str]) -> float:
        """Return fractional adoption rate [0.0, 1.0]."""
        if not workflow_keywords or not step_text:
            return 0.0
        tl = step_text.lower()
        matches = sum(1 for kw in workflow_keywords if kw.lower() in tl)
        return matches / len(workflow_keywords)
