"""Tests for temporal.context.ContextBuilder — fetch_data mark-boundary queries.

Ported from bedrock's memory/tests/test_context.py, using InMemoryStore
instead of Django ORM.
"""

from datetime import date, datetime, timedelta, timezone

from temporal import Message, PeriodType
from temporal.context import ContextBuilder, DefaultMessageFormatter
from temporal.stores import InMemoryStore
from temporal.windows.marks import compute_marks


def _builder(store: InMemoryStore, tz: str = "UTC") -> ContextBuilder:
    return ContextBuilder(
        store=store,
        formatter=DefaultMessageFormatter(tz),
        timezone=tz,
    )


def _seed_five_min(store: InMemoryStore, d: date, hour: int, minute: int, summary: str = "5min"):
    store.upsert_summary(
        PeriodType.FIVE_MINUTE,
        key={"date": d, "hour": hour, "minute": minute},
        summary=summary, message_count=1,
    )


def _seed_hourly(store: InMemoryStore, d: date, hour: int, summary: str = "hourly"):
    store.upsert_summary(
        PeriodType.HOURLY,
        key={"date": d, "hour": hour},
        summary=summary, message_count=1,
    )


def _seed_message(store: InMemoryStore, ts: datetime, content: str = "msg"):
    store.store_messages([Message(role="user", content=content, timestamp=ts)])


class TestMarkBoundaryQueries:
    def test_messages_window_is_after_five_min_cursor(self):
        store = InMemoryStore()
        current_time = datetime(2024, 1, 10, 10, 12, 0, tzinfo=timezone.utc)
        marks = compute_marks(current_time)

        _seed_five_min(store, marks.five_min_cursor.date(), marks.five_min_cursor.hour, 0, "anchor")
        _seed_message(store, marks.five_min_cursor - timedelta(minutes=1), "before-cursor")
        _seed_message(store, marks.five_min_cursor + timedelta(minutes=1), "after-cursor")
        _seed_message(store, current_time + timedelta(minutes=1), "future")

        data = _builder(store).fetch_data(current_time)
        contents = [m.content for m in data.messages]

        assert "after-cursor" in contents
        assert "before-cursor" not in contents
        assert "future" not in contents

    def test_five_minute_summaries_use_marks_range(self):
        store = InMemoryStore()
        current_time = datetime(2024, 1, 10, 10, 33, 0, tzinfo=timezone.utc)
        marks = compute_marks(current_time)

        in_start = marks.five_min_window_start
        in_mid = marks.five_min_window_start + timedelta(minutes=30)
        in_end_minus = marks.five_min_cursor - timedelta(minutes=5)

        _seed_five_min(store, in_start.date(), in_start.hour, in_start.minute, "in-start")
        _seed_five_min(store, in_mid.date(), in_mid.hour, in_mid.minute, "in-mid")
        _seed_five_min(store, in_end_minus.date(), in_end_minus.hour, in_end_minus.minute, "in-end-minus")

        out_before = marks.five_min_window_start - timedelta(minutes=5)
        out_at_end = marks.five_min_cursor
        _seed_five_min(store, out_before.date(), out_before.hour, out_before.minute, "out-before")
        _seed_five_min(store, out_at_end.date(), out_at_end.hour, out_at_end.minute, "out-at-end")

        data = _builder(store).fetch_data(current_time)
        summaries = [s.summary for s in data.five_minute_summaries]

        assert "in-start" in summaries
        assert "in-mid" in summaries
        assert "in-end-minus" in summaries
        assert "out-before" not in summaries
        assert "out-at-end" not in summaries

    def test_hourly_summaries_use_marks_range(self):
        store = InMemoryStore()
        current_time = datetime(2024, 1, 10, 10, 33, 0, tzinfo=timezone.utc)
        marks = compute_marks(current_time)

        in_start = marks.hourly_window_start
        in_mid = marks.hourly_window_start + timedelta(hours=12)
        in_end_minus = marks.five_min_window_start - timedelta(hours=1)

        _seed_hourly(store, in_start.date(), in_start.hour, "in-start")
        _seed_hourly(store, in_mid.date(), in_mid.hour, "in-mid")
        _seed_hourly(store, in_end_minus.date(), in_end_minus.hour, "in-end-minus")

        out_before = marks.hourly_window_start - timedelta(hours=1)
        out_at_end = marks.five_min_window_start
        _seed_hourly(store, out_before.date(), out_before.hour, "out-before")
        _seed_hourly(store, out_at_end.date(), out_at_end.hour, "out-at-end")

        data = _builder(store).fetch_data(current_time)
        summaries = [s.summary for s in data.hourly_summaries]

        assert "in-start" in summaries
        assert "in-mid" in summaries
        assert "in-end-minus" in summaries
        assert "out-before" not in summaries
        assert "out-at-end" not in summaries

    def test_min_resolution_five_minute_excludes_raw_messages(self):
        store = InMemoryStore()
        current_time = datetime(2024, 1, 10, 10, 12, 0, tzinfo=timezone.utc)
        marks = compute_marks(current_time)
        _seed_message(store, marks.five_min_cursor + timedelta(minutes=1), "msg")

        data = _builder(store).fetch_data(current_time, min_resolution=PeriodType.FIVE_MINUTE)
        assert len(data.messages) == 0


class TestPreviousSummaryContext:
    """Ported from bedrock's test_summarizer_previous_summaries.py."""

    def test_five_minute_includes_previous_across_hour_boundary(self):
        store = InMemoryStore()
        _seed_five_min(store, date(2024, 1, 15), 10, 55, "prev-5m")

        builder = _builder(store)
        data = builder.fetch_data(
            datetime(2024, 1, 15, 11, 0, 0, tzinfo=timezone.utc),
            PeriodType.FIVE_MINUTE,
        )
        summaries = [s.summary for s in data.five_minute_summaries]
        assert "prev-5m" in summaries

    def test_hourly_includes_previous_across_day_boundary(self):
        store = InMemoryStore()
        _seed_hourly(store, date(2024, 1, 14), 23, "prev-hour")

        builder = _builder(store)
        data = builder.fetch_data(
            datetime(2024, 1, 15, 0, 0, 0, tzinfo=timezone.utc),
            PeriodType.HOURLY,
        )
        summaries = [s.summary for s in data.hourly_summaries]
        assert "prev-hour" in summaries

    def test_daily_includes_previous_days(self):
        store = InMemoryStore()
        store.upsert_summary(PeriodType.DAILY, key={"date": date(2024, 1, 12)}, summary="d12", message_count=1)
        store.upsert_summary(PeriodType.DAILY, key={"date": date(2024, 1, 13)}, summary="d13", message_count=1)

        builder = _builder(store)
        data = builder.fetch_data(
            datetime(2024, 1, 14, 0, 0, 0, tzinfo=timezone.utc),
            PeriodType.DAILY,
        )
        summaries = [s.summary for s in data.daily_summaries]
        assert "d12" in summaries
        assert "d13" in summaries

    def test_weekly_includes_previous_weeks(self):
        store = InMemoryStore()
        store.upsert_summary(
            PeriodType.WEEKLY,
            key={"week_start_date": date(2024, 1, 7)},
            summary="w-jan7", message_count=1,
        )

        builder = _builder(store)
        data = builder.fetch_data(
            datetime(2024, 1, 15, 0, 0, 0, tzinfo=timezone.utc),
            PeriodType.WEEKLY,
        )
        summaries = [s.summary for s in data.weekly_summaries]
        assert "w-jan7" in summaries

    def test_monthly_includes_previous_months(self):
        store = InMemoryStore()
        store.upsert_summary(
            PeriodType.MONTHLY,
            key={"year": 2023, "month": 12},
            summary="m-dec", message_count=1,
        )

        builder = _builder(store)
        data = builder.fetch_data(
            datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
            PeriodType.MONTHLY,
        )
        summaries = [s.summary for s in data.monthly_summaries]
        assert "m-dec" in summaries
