"""Tests for orin_labs_memory.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 orin_labs_memory import Message, PeriodType
from orin_labs_memory.context import ContextBuilder, DefaultMessageFormatter
from orin_labs_memory.stores import InMemoryStore
from orin_labs_memory.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
