"""
Task Trace Builder — Builds the task chain trace for Inspector.

PRD §4.1 — Task Trace output:
{
    "task_id": "task_001",
    "recalled_memories": [{"memory_id", "memory_layer", "summary"}],
    "injected_memories": [{"memory_id", "token_size", "position"}],
    "adopted_memories":  [{"memory_id"}],
    "workflow_used":     {"workflow_id", "adoption_rate", "summary"},
    "result":            {"success", "execution_time", "summary"}
}

PRD §5.1 — Aggregation logic:
    recalled = filter(events, recalled=True)
    injected = filter(events, injected=True)
    adopted  = filter(events, adopted=True)
    workflow = find workflow feedback for task
    result   = get task outcome

PRD §9 — Performance: task trace build < 50ms
"""

import time
import logging
from typing import List, Dict, Optional
from openmemo.feedback_aggregator.utils import (
    to_iso, safe_div, filter_records, unique_by
)

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


class TaskTraceBuilder:
    """
    Builds a full task trace from raw feedback events.

    Usage:
        builder = TaskTraceBuilder(feedback_store)
        trace = builder.build("task_001")
    """

    def __init__(self, feedback_store):
        self._store = feedback_store

    def build(self, task_id: str) -> dict:
        """
        PRD §5.1 — Build task trace for Inspector.

        Returns structured trace dict.
        """
        t0 = time.time()

        # Load all memory feedback events for this task
        events = self._store.get_task_feedback(task_id)

        # §5.1 filter by signal type
        recalled_events = filter_records(events, recalled=True)
        injected_events = filter_records(events, injected=True)
        adopted_events  = filter_records(events, adopted=True)
        success_events  = filter_records(events, success=True)

        # Deduplicate by memory_id
        recalled_mems = unique_by(recalled_events, "memory_id")
        injected_mems = unique_by(injected_events, "memory_id")
        adopted_mems  = unique_by(adopted_events,  "memory_id")

        # Load workflow feedback for this task
        workflow_used = self._get_workflow(task_id)

        # Task result
        task_success = bool(success_events)
        execution_time = 0.0
        if workflow_used and workflow_used.get("completion_time"):
            execution_time = workflow_used["completion_time"]

        trace = {
            "task_id": task_id,
            "recalled_memories": [
                {
                    "memory_id": e.memory_id,
                    "memory_layer": e.memory_layer or "unknown",
                    "agent_id": e.agent_id,
                    "timestamp": to_iso(e.timestamp),
                }
                for e in recalled_mems
            ],
            "injected_memories": [
                {
                    "memory_id": e.memory_id,
                    "token_size": e.token_size,
                    "position": e.position,
                    "timestamp": to_iso(e.timestamp),
                }
                for e in injected_mems
            ],
            "adopted_memories": [
                {
                    "memory_id": e.memory_id,
                    "agent_id": e.agent_id,
                    "timestamp": to_iso(e.timestamp),
                }
                for e in adopted_mems
            ],
            "workflow_used": workflow_used,
            "result": {
                "success": task_success,
                "execution_time": execution_time,
                "summary": self._result_summary(task_success, adopted_mems),
            },
            "stats": {
                "recalled_count": len(recalled_mems),
                "injected_count": len(injected_mems),
                "adopted_count":  len(adopted_mems),
                "adoption_rate":  safe_div(len(adopted_mems), len(recalled_mems)),
            },
            "build_ms": round((time.time() - t0) * 1000, 1),
        }

        logger.debug("task_trace: task=%s recalled=%d injected=%d adopted=%d ms=%.1f",
                     task_id, len(recalled_mems), len(injected_mems),
                     len(adopted_mems), trace["build_ms"])
        return trace

    def _get_workflow(self, task_id: str) -> Optional[dict]:
        """Find the workflow used in this task from workflow_feedback records."""
        wf_records = self._store.get_task_feedback(task_id)
        # Workflow records have success field and come from workflow_feedback table
        # Try the workflow_feedback store directly
        try:
            wf_rows = self._store.conn.execute(
                "SELECT * FROM workflow_feedback WHERE task_id=? ORDER BY timestamp DESC LIMIT 1",
                (task_id,)
            ).fetchone()
            if wf_rows:
                return {
                    "workflow_id": wf_rows["workflow_id"],
                    "adopted": bool(wf_rows["adopted"]),
                    "success": bool(wf_rows["success"]),
                    "adoption_rate": wf_rows["adoption_rate"],
                    "completion_time": wf_rows["completion_time"],
                    "summary": f"workflow:{wf_rows['workflow_id']}",
                }
        except Exception:
            pass
        return None

    def _result_summary(self, success: bool, adopted_mems: list) -> str:
        if success and adopted_mems:
            return f"Task succeeded; {len(adopted_mems)} memories adopted"
        elif success:
            return "Task succeeded"
        elif adopted_mems:
            return f"Task failed; {len(adopted_mems)} memories were used"
        return "Task failed; no memories adopted"
