"""Integration tests for temporal.service.Memory — the full pipeline."""

from datetime import datetime, timedelta, timezone

import tiktoken

from temporal import Memory, Message, PeriodType
from temporal.core.models import LLMResponse
from temporal.stores import InMemoryStore


class FakeLLM:
    """Deterministic LLM that returns a summary based on the period type keyword in the prompt."""

    def __init__(self):
        self.calls: list[dict] = []

    async def complete(self, prompt: str, model: str, max_tokens: int = 4096) -> LLMResponse:
        self.calls.append({"model": model, "prompt_len": len(prompt)})
        return LLMResponse(
            text="I helped the user with their request.",
            input_tokens=100,
            output_tokens=20,
        )


def _make_big_text(approx_tokens: int) -> str:
    words = [f"word_{i:04d}" for i in range(approx_tokens)]
    text = " ".join(words)
    enc = tiktoken.get_encoding("o200k_base")
    tokens = enc.encode(text)
    if len(tokens) > approx_tokens:
        text = enc.decode(tokens[:approx_tokens])
    return text


class TestMemoryLog:
    def test_log_assigns_timestamps(self):
        store = InMemoryStore()
        memory = Memory(store=store, llm=FakeLLM())

        ts = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        memory.log([
            Message(role="user", content="First"),
            Message(role="assistant", content="Second"),
        ], timestamp=ts)

        msgs = store.query_messages()
        assert len(msgs) == 2
        assert msgs[0].timestamp < msgs[1].timestamp
        assert msgs[0].content == "First"

    def test_log_preserves_existing_timestamps(self):
        store = InMemoryStore()
        memory = Memory(store=store, llm=FakeLLM())

        ts = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        custom_ts = datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
        memory.log([
            Message(role="user", content="Explicit", timestamp=custom_ts),
        ], timestamp=ts)

        msgs = store.query_messages()
        assert msgs[0].timestamp == custom_ts


class TestMemorySummarize:
    def test_summarize_creates_five_minute_summary(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        past = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        memory.log([
            Message(role="user", content="What is the weather?"),
            Message(role="assistant", content="It's sunny."),
        ], timestamp=past)

        now = datetime(2026, 3, 2, 10, 10, 0, tzinfo=timezone.utc)
        result = memory.summarize(current_time=now)

        assert result["five_minute_created"] >= 1
        assert result["llm_usage"]["llm_calls"] >= 1
        assert store.has_summaries(PeriodType.FIVE_MINUTE)

    def test_summarize_idempotent(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        past = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        memory.log([Message(role="user", content="Hello")], timestamp=past)

        now = datetime(2026, 3, 2, 10, 10, 0, tzinfo=timezone.utc)
        r1 = memory.summarize(current_time=now)
        calls_after_first = len(llm.calls)
        r2 = memory.summarize(current_time=now)

        assert r2["five_minute_created"] == 0
        assert len(llm.calls) == calls_after_first


class TestMemoryRecall:
    def test_recall_returns_context_and_messages(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        past = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        memory.log([
            Message(role="user", content="What is 2+2?"),
            Message(role="assistant", content="4"),
        ], timestamp=past)

        now = datetime(2026, 3, 2, 10, 10, 0, tzinfo=timezone.utc)
        memory.summarize(current_time=now)

        system_text, conversation = memory.recall(current_time=now)

        assert "5-MINUTE SUMMARIES" in system_text
        assert isinstance(conversation, list)

    def test_recall_two_step_fetch_render(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        past = datetime(2026, 3, 2, 10, 0, 0, tzinfo=timezone.utc)
        memory.log([Message(role="user", content="Hi")], timestamp=past)
        now = datetime(2026, 3, 2, 10, 10, 0, tzinfo=timezone.utc)
        memory.summarize(current_time=now)

        data = memory.fetch(current_time=now)
        assert len(data.five_minute_summaries) >= 1

        msgs = list(data.messages)
        data.messages = []
        text = memory.render(data)

        assert "5-MINUTE SUMMARIES" in text
        assert isinstance(msgs, list)


class TestMemoryTrimming:
    def test_recall_trims_large_messages(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        big_content = _make_big_text(10_000)
        now = datetime(2026, 3, 2, 10, 12, 0, tzinfo=timezone.utc)
        base = now - timedelta(minutes=1)

        store.store_messages([
            Message(role="user", content="Short question", timestamp=base),
            Message(role="assistant", content=big_content, timestamp=base + timedelta(seconds=1)),
            Message(role="user", content="Follow up?", timestamp=base + timedelta(seconds=2)),
        ])

        store.upsert_summary(
            PeriodType.FIVE_MINUTE,
            key={"date": now.date(), "hour": 10, "minute": 0},
            summary="anchor",
            message_count=1,
        )

        system_text, conversation = memory.recall(current_time=now, max_tokens=2000)

        enc = tiktoken.get_encoding("o200k_base")
        total = len(enc.encode(system_text))
        for msg in conversation:
            total += len(enc.encode(msg.plain_text()))

        assert total < 10_000, f"Expected trimming, got {total} tokens"

    def test_recall_no_trimming_when_small(self):
        store = InMemoryStore()
        llm = FakeLLM()
        memory = Memory(store=store, llm=llm, timezone="UTC")

        now = datetime(2026, 3, 2, 10, 12, 0, tzinfo=timezone.utc)
        base = now - timedelta(minutes=1)
        store.store_messages([
            Message(role="user", content="Hello!", timestamp=base),
            Message(role="assistant", content="Hi there!", timestamp=base + timedelta(seconds=1)),
        ])
        store.upsert_summary(
            PeriodType.FIVE_MINUTE,
            key={"date": now.date(), "hour": 10, "minute": 0},
            summary="anchor",
            message_count=1,
        )

        system_text, conversation = memory.recall(current_time=now, max_tokens=150_000)

        for msg in conversation:
            assert "[trimmed to fit context window]" not in msg.plain_text()
