from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING

import tiktoken

from temporal.context.builder import ContextBuilder
from temporal.context.compression import DefaultMessageFormatter
from temporal.core.models import MemoryData, Message
from temporal.core.types import PeriodType
from temporal.core.utils import force_timezone
from temporal.pipeline.summarizer import Summarizer

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

logger = logging.getLogger(__name__)


class Memory:
    """Main interface for temporal memory.

    Provides hierarchical summarization of conversations across time scales
    (5-minute, hourly, daily, weekly, monthly) with context rendering and
    optional associative recall.

    Example::

        from temporal import Memory, Message
        from temporal.stores import InMemoryStore
        from temporal.llm import OpenAIClient

        memory = Memory(
            store=InMemoryStore(),
            llm=OpenAIClient(),
        )

        memory.log([
            Message(role="user", content="Hello!"),
            Message(role="assistant", content="Hi there!"),
        ])

        memory.summarize()
        context = memory.recall()
    """

    def __init__(
        self,
        store: MemoryStore,
        llm: LLMClient,
        *,
        timezone: str = "UTC",
        summarizer_model: str = "gpt-4.1-mini",
        formatter: MessageFormatter | None = None,
        time_offset: int = 0,
    ) -> None:
        self.store = store
        self.llm = llm
        self.timezone = timezone
        self.summarizer_model = summarizer_model
        self.time_offset = time_offset

        self._formatter = formatter or DefaultMessageFormatter(timezone)
        self._summarizer = Summarizer(
            store=store,
            llm=llm,
            timezone=timezone,
            model=summarizer_model,
            formatter=self._formatter,
        )
        self._context_builder = ContextBuilder(
            store=store,
            formatter=self._formatter,
            timezone=timezone,
            time_offset=time_offset,
        )

    def log(
        self,
        messages: list[Message],
        timestamp: datetime | None = None,
    ) -> None:
        """Log messages to memory.

        Messages without a ``timestamp`` field will be assigned one based
        on the provided *timestamp* (or ``now()``), incrementing by 1ms
        per message to preserve ordering.
        """
        if timestamp is None:
            timestamp = datetime.now(timezone.utc)

        timestamp = force_timezone(timestamp, self.timezone)
        timestamp = timestamp + timedelta(minutes=self.time_offset)

        stamped: list[Message] = []
        for idx, msg in enumerate(messages):
            if msg.timestamp is None:
                ts = timestamp + timedelta(milliseconds=idx)
                stamped.append(msg.model_copy(update={"timestamp": ts}))
            else:
                stamped.append(msg)

        self.store.store_messages(stamped)

    def summarize(self, current_time: datetime | None = None) -> dict:
        """Run hierarchical summarization across all time levels.

        Returns a dict with counts of summaries created per level and LLM usage.
        """
        if current_time is None:
            current_time = datetime.now(timezone.utc)
        return self._summarizer.update_all(current_time)

    def fetch(
        self,
        current_time: datetime | None = None,
        min_resolution: PeriodType | None = None,
    ) -> MemoryData:
        """Fetch all memory data (summaries + messages) from the store.

        Returns a ``MemoryData`` bundle the caller can inspect and modify
        before passing to ``render()``.
        """
        if current_time is None:
            current_time = datetime.now(timezone.utc)

        current_time = force_timezone(current_time, self.timezone)
        current_time = current_time + timedelta(minutes=self.time_offset)
        return self._context_builder.fetch_data(current_time, min_resolution)

    def render(self, data: MemoryData, max_tokens: int = 50_000) -> str:
        """Render a ``MemoryData`` bundle into a single context string."""
        return self._context_builder.render(data, max_tokens)

    def recall(
        self,
        current_time: datetime | None = None,
        max_tokens: int = 150_000,
    ) -> tuple[str, list[Message]]:
        """Build a complete memory context and conversation for an LLM request.

        Returns ``(system_context, conversation_messages)`` where:
        - ``system_context`` is a rendered string of all summaries
        - ``conversation_messages`` is a list of recent Messages for the
          conversation history

        The caller typically puts ``system_context`` in the system prompt
        and ``conversation_messages`` in the messages array.
        """
        data = self.fetch(current_time=current_time)
        recent_messages = list(data.messages)
        data.messages = []
        system_text = self.render(data)

        self._trim_if_needed(system_text, recent_messages, max_tokens)
        return system_text, recent_messages

    def _trim_if_needed(
        self,
        system_text: str,
        messages: list[Message],
        max_total_tokens: int,
    ) -> None:
        """Trim the longest message content blocks to fit within budget.

        Modifies *messages* in place.
        """
        if not messages:
            return

        enc = tiktoken.get_encoding("o200k_base")
        system_tokens = len(enc.encode(system_text)) if system_text else 0

        msg_tokens = []
        total_msg_tokens = 0
        for idx, msg in enumerate(messages):
            text = msg.plain_text()
            tokens = len(enc.encode(text)) if text else 0
            msg_tokens.append((tokens, idx))
            total_msg_tokens += tokens

        total = system_tokens + total_msg_tokens
        if total <= max_total_tokens:
            return

        overflow = total - max_total_tokens
        logger.warning(
            "Memory + conversation tokens (%d) exceed budget (%d) by %d. Trimming.",
            total, max_total_tokens, overflow,
        )

        MIN_PRESERVED = 100
        trimmable = sorted(
            [(t, i) for t, i in msg_tokens if t > MIN_PRESERVED],
            key=lambda x: x[0],
            reverse=True,
        )

        for tokens, idx in trimmable:
            if overflow <= 0:
                break

            msg = messages[idx]
            if isinstance(msg.content, str):
                target = max(MIN_PRESERVED, tokens - overflow)
                trimmed = self._truncate(msg.content, target, enc)
                reduction = tokens - len(enc.encode(trimmed))
                messages[idx] = msg.model_copy(update={"content": trimmed})
                overflow -= reduction

    @staticmethod
    def _truncate(text: str, max_tokens: int, enc) -> str:
        if not text:
            return text
        current = enc.encode(text)
        if len(current) <= max_tokens:
            return text
        notice = "\n\n... [trimmed to fit context window] ..."
        notice_tokens = len(enc.encode(notice))
        keep = max(0, max_tokens - notice_tokens)
        return enc.decode(current[:keep]) + notice
