from __future__ import annotations

import threading
from datetime import date, datetime
from typing import Any

from temporal.core.models import Message, SummaryRecord
from temporal.core.types import PeriodType


class InMemoryStore:
    """Dict-backed memory store. Good for testing, prototyping, and short-lived sessions."""

    def __init__(self) -> None:
        self._messages: list[Message] = []
        self._summaries: dict[PeriodType, list[SummaryRecord]] = {p: [] for p in PeriodType}
        self._lock = threading.Lock()

    def store_messages(self, messages: list[Message]) -> None:
        with self._lock:
            self._messages.extend(messages)

    def query_messages(
        self,
        *,
        after: datetime | None = None,
        before: datetime | None = None,
        limit: int | None = None,
    ) -> list[Message]:
        with self._lock:
            result = []
            for msg in self._messages:
                ts = msg.timestamp
                if ts is None:
                    continue
                if after is not None and ts <= after:
                    continue
                if before is not None and ts > before:
                    continue
                result.append(msg)
            result.sort(key=lambda m: m.timestamp or datetime.min)
            if limit is not None:
                result = result[:limit]
            return result

    def get_raw_message_data(
        self,
        *,
        after: datetime | None = None,
        before: datetime | None = None,
    ) -> list[tuple[datetime, Any]]:
        with self._lock:
            result = []
            for msg in self._messages:
                ts = msg.timestamp
                if ts is None:
                    continue
                if after is not None and ts <= after:
                    continue
                if before is not None and ts >= before:
                    continue
                result.append((ts, msg.to_dict()))
            result.sort(key=lambda x: x[0])
            return result

    def has_summaries(self, period: PeriodType) -> bool:
        with self._lock:
            return len(self._summaries[period]) > 0

    def get_summary_keys(self, period: PeriodType) -> set[tuple]:
        with self._lock:
            return {self._key_for(period, s) for s in self._summaries[period]}

    def query_summaries(
        self,
        period: PeriodType,
        *,
        start: datetime | None = None,
        end: datetime | None = None,
        limit: int | None = None,
    ) -> list[SummaryRecord]:
        with self._lock:
            result = []
            for s in self._summaries[period]:
                if not self._in_range(period, s, start, end):
                    continue
                result.append(s)
            result.sort(key=lambda s: self._sort_key(period, s))
            if limit is not None:
                result = result[:limit]
            return result

    def get_all_summaries(self, period: PeriodType) -> list[SummaryRecord]:
        with self._lock:
            result = list(self._summaries[period])
            result.sort(key=lambda s: self._sort_key(period, s))
            return result

    def get_summaries_before(self, period: PeriodType, before: date) -> list[SummaryRecord]:
        with self._lock:
            result = []
            for s in self._summaries[period]:
                if self._date_of(period, s) < before:
                    result.append(s)
            result.sort(key=lambda s: self._sort_key(period, s))
            return result

    def upsert_summary(
        self,
        period: PeriodType,
        key: dict[str, Any],
        summary: str,
        message_count: int,
    ) -> bool:
        with self._lock:
            record = SummaryRecord(summary=summary, message_count=message_count, **key)
            new_key = self._key_for(period, record)

            for i, existing in enumerate(self._summaries[period]):
                if self._key_for(period, existing) == new_key:
                    self._summaries[period][i] = record
                    return False  # updated, not created

            self._summaries[period].append(record)
            return True  # created

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _key_for(period: PeriodType, s: SummaryRecord) -> tuple:
        if period == PeriodType.FIVE_MINUTE:
            return (s.date, s.hour, s.minute)
        if period == PeriodType.HOURLY:
            return (s.date, s.hour)
        if period == PeriodType.DAILY:
            return (s.date,)
        if period == PeriodType.WEEKLY:
            return (s.week_start_date,)
        if period == PeriodType.MONTHLY:
            return (s.year, s.month)
        return ()

    @staticmethod
    def _sort_key(period: PeriodType, s: SummaryRecord) -> tuple:
        if period == PeriodType.FIVE_MINUTE:
            return (s.date or date.min, s.hour or 0, s.minute or 0)
        if period == PeriodType.HOURLY:
            return (s.date or date.min, s.hour or 0)
        if period == PeriodType.DAILY:
            return (s.date or date.min,)
        if period == PeriodType.WEEKLY:
            return (s.week_start_date or date.min,)
        if period == PeriodType.MONTHLY:
            return (s.year or 0, s.month or 0)
        return ()

    @staticmethod
    def _date_of(period: PeriodType, s: SummaryRecord) -> date:
        if period in (PeriodType.FIVE_MINUTE, PeriodType.HOURLY, PeriodType.DAILY):
            return s.date or date.min
        if period == PeriodType.WEEKLY:
            return s.week_start_date or date.min
        if period == PeriodType.MONTHLY:
            return date(s.year or 1, s.month or 1, 1)
        return date.min

    @staticmethod
    def _in_range(
        period: PeriodType,
        s: SummaryRecord,
        start: datetime | None,
        end: datetime | None,
    ) -> bool:
        if period == PeriodType.FIVE_MINUTE:
            key = (s.date, s.hour or 0, s.minute or 0)
            if start:
                s_key = (start.date(), start.hour, start.minute)
                if key < s_key:
                    return False
            if end:
                e_key = (end.date(), end.hour, end.minute)
                if key >= e_key:
                    return False
            return True

        if period == PeriodType.HOURLY:
            key = (s.date, s.hour or 0)
            if start:
                s_key = (start.date(), start.hour)
                if key < s_key:
                    return False
            if end:
                e_key = (end.date(), end.hour)
                if key >= e_key:
                    return False
            return True

        if period == PeriodType.DAILY:
            d = s.date or date.min
            if start and d < start.date():
                return False
            if end and d >= end.date():
                return False
            return True

        if period == PeriodType.WEEKLY:
            d = s.week_start_date or date.min
            if start and d < start.date():
                return False
            if end and d >= end.date():
                return False
            return True

        if period == PeriodType.MONTHLY:
            key = (s.year or 0, s.month or 0)
            if end:
                e_key = (end.year, end.month)
                if key >= e_key:
                    return False
            return True

        return True
