from __future__ import annotations

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

import tiktoken

from temporal.core.models import MemoryData, SummaryRecord
from temporal.core.types import PeriodType
from temporal.core.utils import force_timezone
from temporal.windows.marks import compute_windows

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

logger = logging.getLogger(__name__)

RESOLUTION_LEVEL: dict[PeriodType, int] = {
    PeriodType.FIVE_MINUTE: 0,
    PeriodType.HOURLY: 1,
    PeriodType.DAILY: 2,
    PeriodType.WEEKLY: 3,
    PeriodType.MONTHLY: 4,
}


class ContextBuilder:
    """Constructs a time-aware memory context using stratified summaries."""

    def __init__(
        self,
        store: MemoryStore,
        formatter: MessageFormatter,
        timezone: str,
        time_offset: int = 0,
    ) -> None:
        self.store = store
        self.formatter = formatter
        self.timezone = timezone
        self.time_offset = time_offset

    def fetch_data(
        self,
        current_time: datetime,
        min_resolution: PeriodType | None = None,
    ) -> MemoryData:
        """Fetch summaries/messages strictly before *current_time* using mark boundaries."""
        local_now = force_timezone(current_time, self.timezone)
        windows = compute_windows(local_now, min_resolution)
        min_level = RESOLUTION_LEVEL.get(min_resolution, 0) if min_resolution else -1

        from temporal.core.models import Message

        messages: list[Message] = []
        if min_resolution is None:
            assert windows.message_start is not None and windows.message_end is not None
            lower_bound = windows.message_start
            if not self.store.has_summaries(PeriodType.FIVE_MINUTE):
                lower_bound = windows.message_end.replace(minute=0, second=0, microsecond=0)
            messages = self.store.query_messages(
                after=lower_bound,
                before=windows.message_end,
                limit=200,
            )

        five_minute_summaries: list[SummaryRecord] = []
        if min_level <= RESOLUTION_LEVEL[PeriodType.FIVE_MINUTE]:
            assert windows.five_min_start is not None and windows.five_min_end is not None
            five_minute_summaries = self.store.query_summaries(
                PeriodType.FIVE_MINUTE,
                start=windows.five_min_start,
                end=windows.five_min_end,
            )

        hourly_summaries: list[SummaryRecord] = []
        if min_level <= RESOLUTION_LEVEL[PeriodType.HOURLY]:
            assert windows.hourly_start is not None and windows.hourly_end is not None
            hourly_summaries = self.store.query_summaries(
                PeriodType.HOURLY,
                start=windows.hourly_start,
                end=windows.hourly_end,
            )

        daily_summaries: list[SummaryRecord] = []
        if min_level <= RESOLUTION_LEVEL[PeriodType.DAILY]:
            assert windows.daily_start is not None and windows.daily_end is not None
            daily_summaries = self.store.query_summaries(
                PeriodType.DAILY,
                start=windows.daily_start,
                end=windows.daily_end,
            )

        weekly_summaries: list[SummaryRecord] = []
        if min_level <= RESOLUTION_LEVEL[PeriodType.WEEKLY]:
            assert windows.weekly_start is not None and windows.weekly_end is not None
            weekly_summaries = self.store.query_summaries(
                PeriodType.WEEKLY,
                start=windows.weekly_start,
                end=windows.weekly_end,
            )

        monthly_summaries: list[SummaryRecord] = []
        if min_level <= RESOLUTION_LEVEL[PeriodType.MONTHLY]:
            monthly_summaries = self.store.query_summaries(
                PeriodType.MONTHLY,
                end=windows.monthly_end,
                limit=24,
            )

        return MemoryData(
            monthly_summaries=monthly_summaries,
            weekly_summaries=weekly_summaries,
            daily_summaries=daily_summaries,
            hourly_summaries=hourly_summaries,
            five_minute_summaries=five_minute_summaries,
            messages=messages,
            last_summarized_time=windows.message_start or windows.five_min_end,
        )

    def render(self, data: MemoryData, max_tokens: int = 50_000) -> str:
        """Render a MemoryData bundle into a single context string."""
        enc = tiktoken.get_encoding("o200k_base")
        return self._render_sections(data, max_tokens, enc)

    def _render_sections(self, data: MemoryData, max_tokens: int, enc) -> str:
        parts: list[str] = []

        if data.monthly_summaries:
            parts.append("=== MONTHLY SUMMARIES ===")
            for s in data.monthly_summaries:
                parts.append(f"\n{s.year}-{(s.month or 0):02d}: {s.summary}")

        if data.weekly_summaries:
            parts.append("\n\n=== RECENT WEEKLY SUMMARIES ===")
            for s in data.weekly_summaries:
                parts.append(f"\nWeek of {s.week_start_date}: {s.summary}")

        if data.daily_summaries:
            parts.append("\n\n=== RECENT DAILY SUMMARIES ===")
            for s in data.daily_summaries:
                parts.append(f"\n{s.date}: {s.summary}")

        if data.hourly_summaries:
            parts.append("=== HOURLY SUMMARIES ===")
            for s in data.hourly_summaries:
                dt = datetime.combine(s.date, time(s.hour or 0, 0, 0))
                dt = force_timezone(dt, self.timezone) + timedelta(minutes=self.time_offset)
                parts.append(f"\n{dt.strftime('%H:%M:%S')}: {s.summary}")

        if data.five_minute_summaries:
            parts.append("\n\n=== 5-MINUTE SUMMARIES ===")
            for s in data.five_minute_summaries:
                start_dt = datetime.combine(s.date, time(s.hour or 0, s.minute or 0, 0))
                start_dt = force_timezone(start_dt, self.timezone)
                end_dt = start_dt + timedelta(minutes=5)
                if self.time_offset:
                    start_dt += timedelta(minutes=self.time_offset)
                    end_dt += timedelta(minutes=self.time_offset)
                parts.append(
                    f"\n{start_dt.strftime('%H:%M')}–{end_dt.strftime('%H:%M')}: {s.summary}"
                )

        if data.messages:
            current_text = "\n".join(parts)
            current_tokens = len(enc.encode(current_text)) if current_text else 0
            remaining_tokens = max_tokens - current_tokens

            if remaining_tokens > 0:
                formatted = self.formatter.apply_compression_strategies(
                    data.messages, remaining_tokens, enc
                )
                if formatted:
                    parts.append("\n\n=== MOST RECENT MESSAGES ===")
                    parts.extend(formatted)

        return "\n".join(parts) if parts else ""
