from __future__ import annotations

import asyncio
import concurrent.futures
import contextvars
import logging
import textwrap
import time as time_mod
from dataclasses import dataclass, field
from datetime import datetime
from typing import TYPE_CHECKING

import tiktoken

from temporal.core.models import ContentBlock, LLMResponse, Message, SummaryRecord
from temporal.core.types import PeriodType

if TYPE_CHECKING:
    from temporal.core.protocols import LLMClient, MemoryStore

logger = logging.getLogger(__name__)

SITUATION_MODEL = "gpt-4.1-nano"
RELEVANCE_MODEL = "gpt-4.1-nano"
AGGREGATION_MODEL = "gpt-4.1-nano"

MAX_SITUATION_TOKENS = 2000
MAX_MESSAGES_FOR_SITUATION = 10
RELEVANCE_TIMEOUT = 30
MAX_RELEVANT_FRAGMENTS = 20


@dataclass
class AssociativeResult:
    """Result of an associative memory query."""

    memory_text: str | None = None
    situation: str | None = None
    input_tokens: int = 0
    output_tokens: int = 0
    llm_calls: int = 0
    metadata: dict = field(default_factory=dict)


class AssociativeMemory:
    """Automatic associative memory pipeline.

    Three-step process:
    1. Extract a detailed situation description from memory context + recent messages
    2. Score all 5-minute summaries for relevance
    3. Aggregate top-N relevant summaries into a consolidated memory block
    """

    def __init__(
        self,
        store: MemoryStore,
        llm: LLMClient,
        *,
        fetch_context: callable | None = None,
        system_prompt: str = "",
        situation_model: str = SITUATION_MODEL,
        relevance_model: str = RELEVANCE_MODEL,
        aggregation_model: str = AGGREGATION_MODEL,
    ) -> None:
        self.store = store
        self.llm = llm
        self.fetch_context = fetch_context
        self.system_prompt = system_prompt
        self.situation_model = situation_model
        self.relevance_model = relevance_model
        self.aggregation_model = aggregation_model

    def query_sync(self, current_time: datetime) -> AssociativeResult:
        """Synchronous wrapper."""
        ctx = contextvars.copy_context()
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            return pool.submit(
                ctx.run,
                asyncio.run,
                self.query(current_time),
            ).result(timeout=300)

    async def query(self, current_time: datetime) -> AssociativeResult:
        """Run the full associative memory pipeline."""
        total_start = time_mod.perf_counter()
        metadata: dict = {
            "situation_model": self.situation_model,
            "relevance_model": self.relevance_model,
            "aggregation_model": self.aggregation_model,
        }

        summaries = self.store.get_all_summaries(PeriodType.FIVE_MINUTE)
        metadata["summaries_total"] = len(summaries)

        if not summaries:
            return AssociativeResult(metadata=metadata)

        # Build situation inputs
        memory_context = ""
        recent_messages_text = ""
        if self.fetch_context:
            memory_context, recent_messages_text = self.fetch_context(current_time)

        situation, sit_usage = await self._extract_situation(memory_context, recent_messages_text)
        metadata["situation_chars"] = len(situation or "")
        if not situation:
            return AssociativeResult(
                input_tokens=sit_usage.input_tokens,
                output_tokens=sit_usage.output_tokens,
                llm_calls=1,
                metadata=metadata,
            )

        # Score all summaries for relevance
        task_to_summary: dict[asyncio.Task, SummaryRecord] = {}
        for s in summaries:
            task = asyncio.create_task(self._score_relevance(situation, s))
            task_to_summary[task] = s

        relevance_start = time_mod.perf_counter()
        done, pending_tasks = await asyncio.wait(
            task_to_summary.keys(), timeout=RELEVANCE_TIMEOUT
        )
        for task in pending_tasks:
            task.cancel()

        metadata["relevance_ms"] = round((time_mod.perf_counter() - relevance_start) * 1000, 2)
        metadata["summaries_scored"] = len(done)
        metadata["summaries_timed_out"] = len(pending_tasks)

        relevance_in = 0
        relevance_out = 0
        scored: list[tuple[int, SummaryRecord]] = []
        for task in done:
            exc = task.exception()
            if exc is not None:
                logger.warning("Associative memory score failed: %s", exc)
                continue
            result = task.result()
            if result is not None:
                score, usage = result
                relevance_in += usage.input_tokens
                relevance_out += usage.output_tokens
                if score >= 20:
                    scored.append((score, task_to_summary[task]))

        if not scored:
            total_in = sit_usage.input_tokens + relevance_in
            total_out = sit_usage.output_tokens + relevance_out
            return AssociativeResult(
                situation=situation,
                input_tokens=total_in,
                output_tokens=total_out,
                llm_calls=1 + len(done),
                metadata=metadata,
            )

        scored.sort(key=lambda x: x[0], reverse=True)
        top = scored[:MAX_RELEVANT_FRAGMENTS]
        top_chrono = sorted(top, key=lambda x: (x[1].date, x[1].hour or 0, x[1].minute or 0))
        fragments = [
            (f"{s.date} {(s.hour or 0):02d}:{(s.minute or 0):02d}", s.summary)
            for _, s in top_chrono
        ]

        memory_text, agg_usage = await self._aggregate(situation, fragments)
        metadata["relevant_count"] = len(top)
        metadata["memory_output_chars"] = len(memory_text or "")
        metadata["total_ms"] = round((time_mod.perf_counter() - total_start) * 1000, 2)

        total_in = sit_usage.input_tokens + relevance_in + agg_usage.input_tokens
        total_out = sit_usage.output_tokens + relevance_out + agg_usage.output_tokens
        total_calls = 1 + len(done) + 1

        return AssociativeResult(
            memory_text=memory_text,
            situation=situation,
            input_tokens=total_in,
            output_tokens=total_out,
            llm_calls=total_calls,
            metadata=metadata,
        )

    async def _extract_situation(
        self, memory_context: str, recent_messages_text: str
    ) -> tuple[str, LLMResponse]:
        enc = tiktoken.get_encoding("o200k_base")
        trimmed_context = _truncate(memory_context, MAX_SITUATION_TOKENS, enc)
        trimmed_messages = _truncate(recent_messages_text, MAX_SITUATION_TOKENS, enc)

        prompt = textwrap.dedent(f"""\
            You are analyzing an AI agent's current situation to trigger associative memory recall.

            Agent's system prompt:
            {self.system_prompt}

            Recent memory context (summaries of past activity):
            {trimmed_context}

            Recent conversation (most recent messages):
            {trimmed_messages}

            Describe the agent's current situation in 2-3 sentences with HIGH SPECIFICITY.""")

        resp = await self.llm.complete(prompt, self.situation_model)
        return resp.text, resp

    async def _score_relevance(
        self, situation: str, summary: SummaryRecord
    ) -> tuple[int, LLMResponse] | None:
        timestamp_str = f"{summary.date} {(summary.hour or 0):02d}:{(summary.minute or 0):02d}"

        prompt = (
            "You are scoring how relevant a past memory is to an AI agent's "
            "current situation.\n\n"
            f"Current situation:\n{situation}\n\n"
            f"Memory from {timestamp_str}:\n{summary.summary}\n\n"
            "Score 0-100. Only score above 50 if the memory contains specific "
            "facts, preferences, decisions, or instructions the agent could NOT "
            "infer from its current context alone. Do NOT score high just because "
            "the memory mentions the same person or topic.\n\n"
            "Reply with ONLY the integer score, nothing else."
        )

        resp = await self.llm.complete(prompt, self.relevance_model)
        try:
            score = int(resp.text.strip())
        except ValueError:
            return None
        return (max(0, min(100, score)), resp)

    async def _aggregate(
        self, situation: str, fragments: list[tuple[str, str]]
    ) -> tuple[str, LLMResponse]:
        text_block = "\n\n".join(f"[{ts}] {text}" for ts, text in fragments)

        prompt = (
            "You are consolidating relevant memories for an AI agent.\n\n"
            f"Current situation:\n{situation}\n\n"
            "The following memory summaries have been scored as most relevant, "
            "in chronological order. More recent entries may overwrite or "
            "supersede older ones on the same topics:\n\n"
            f"{text_block}\n\n"
            "Output a concise, consolidated summary of what the agent should "
            "keep in mind for their current situation. Extract only the specific "
            "facts, preferences, decisions, and instructions that matter right "
            "now. Where information conflicts, prefer the more recent entry. "
            "Be direct and factual."
        )

        resp = await self.llm.complete(prompt, self.aggregation_model)
        text = resp.text.strip() if resp.text else "Failed to consolidate memories."
        return text, resp


def _truncate(text: str, max_tokens: int, enc) -> str:
    if not text:
        return text
    tokens = enc.encode(text)
    if len(tokens) <= max_tokens:
        return text
    return enc.decode(tokens[:max_tokens]) + "\n... [truncated]"
