from __future__ import annotations

import asyncio
import concurrent.futures
import json
import logging
from datetime import date, datetime, time, timedelta
from typing import TYPE_CHECKING, Any

import tiktoken

from temporal.context.builder import ContextBuilder
from temporal.context.compression import DefaultMessageFormatter
from temporal.core.models import LLMResponse, SummaryRecord
from temporal.core.types import PERIOD_META, PeriodType
from temporal.core.utils import force_timezone
from temporal.windows.bucketing import hour_start, last_completed_5m_end
from temporal.windows.marks import week_start_sunday

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

logger = logging.getLogger(__name__)


class Summarizer:
    """Hierarchical summarization pipeline: 5m -> hourly -> daily -> weekly -> monthly."""

    def __init__(
        self,
        store: MemoryStore,
        llm: LLMClient,
        timezone: str,
        model: str = "gpt-4.1-mini",
        formatter: MessageFormatter | None = None,
    ) -> None:
        self.store = store
        self.llm = llm
        self.timezone = timezone
        self.model = model
        self.formatter = formatter or DefaultMessageFormatter(timezone)
        self.total_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0, "llm_calls": 0}

    def update_all(self, current_time: datetime) -> dict[str, Any]:
        """Run all summarization levels synchronously. Returns stats dict."""
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
            result = pool.submit(
                asyncio.run,
                self._update_all_async(current_time),
            ).result(timeout=900)
        return result

    async def _update_all_async(self, current_time: datetime) -> dict[str, Any]:
        self.total_usage = {"input_tokens": 0, "output_tokens": 0, "llm_calls": 0}

        updated_5m = await self._update_five_minute(current_time)
        updated_hour = await self._update_hourly(current_time)
        updated_day = await self._update_daily(current_time)
        updated_week = await self._update_weekly(current_time)
        updated_month = await self._update_monthly(current_time)

        return {
            "five_minute_created": updated_5m,
            "hourly_created": updated_hour,
            "daily_created": updated_day,
            "weekly_created": updated_week,
            "monthly_created": updated_month,
            "llm_usage": dict(self.total_usage),
        }

    # ------------------------------------------------------------------
    # Five-minute summaries
    # ------------------------------------------------------------------

    async def _update_five_minute(self, current_time: datetime) -> int:
        tz_now = force_timezone(current_time, self.timezone)
        cutoff = last_completed_5m_end(tz_now)

        existing_keys = self.store.get_summary_keys(PeriodType.FIVE_MINUTE)

        after_filter = None
        if existing_keys:
            latest_key = max(existing_keys)
            after_filter = force_timezone(
                datetime.combine(latest_key[0], time(latest_key[1], latest_key[2])),
                self.timezone,
            )

        raw_messages = self.store.get_raw_message_data(after=after_filter, before=cutoff)

        buckets: dict[datetime, list[dict]] = {}
        for timestamp, message_data in raw_messages:
            local_time = force_timezone(timestamp, self.timezone)
            bucket_start = last_completed_5m_end(local_time)
            bucket_end = bucket_start + timedelta(minutes=5)
            if bucket_end > tz_now:
                continue
            buckets.setdefault(bucket_start, []).append(
                message_data if isinstance(message_data, dict) else {"content": str(message_data)}
            )

        pending = []
        for bucket_start, bucket_messages in buckets.items():
            key = (bucket_start.date(), bucket_start.hour, bucket_start.minute)
            if key not in existing_keys:
                pending.append((key, bucket_messages))

        if not pending:
            return 0

        existing_memory = await self._get_context_for_summary(tz_now, PeriodType.FIVE_MINUTE)

        created = 0
        for (date_key, hour_key, minute_key), bucket_messages in sorted(pending):
            content = self._prepare_content(bucket_messages)
            summary_text = await self._create_summary(
                content=content,
                period_type=PeriodType.FIVE_MINUTE,
                existing_memory=existing_memory,
            )
            if not summary_text or not summary_text.strip():
                continue

            was_created = self.store.upsert_summary(
                PeriodType.FIVE_MINUTE,
                key={"date": date_key, "hour": hour_key, "minute": minute_key},
                summary=summary_text,
                message_count=len(bucket_messages),
            )
            if was_created:
                created += 1

        return created

    # ------------------------------------------------------------------
    # Hourly summaries
    # ------------------------------------------------------------------

    async def _update_hourly(self, current_time: datetime) -> int:
        tz_now = force_timezone(current_time, self.timezone)
        cutoff = hour_start(tz_now)

        five_min_summaries = self.store.get_all_summaries(PeriodType.FIVE_MINUTE)

        hourly_groups: dict[tuple[date, int], list[SummaryRecord]] = {}
        for s in five_min_summaries:
            key = (s.date, s.hour or 0)
            hourly_groups.setdefault(key, []).append(s)

        existing_keys = self.store.get_summary_keys(PeriodType.HOURLY)

        pending = []
        for (summary_date, hr), summaries_5m in hourly_groups.items():
            if summary_date == cutoff.date() and hr >= cutoff.hour:
                continue
            if (summary_date, hr) not in existing_keys and summaries_5m:
                parts = [f"[{(s.hour or 0):02d}:{(s.minute or 0):02d}] {s.summary}" for s in summaries_5m]
                combined = "\n\n".join(parts)
                total = sum(s.message_count for s in summaries_5m)
                pending.append(((summary_date, hr), combined, total))

        if not pending:
            return 0

        existing_memory = await self._get_context_for_summary(
            force_timezone(current_time, self.timezone), PeriodType.HOURLY
        )

        created = 0
        for (summary_date, hr), content, total_messages in sorted(pending):
            content = self._trim_content(content)
            summary_text = await self._create_summary(
                content=content,
                period_type=PeriodType.HOURLY,
                existing_memory=existing_memory,
            )
            if not summary_text or not summary_text.strip():
                continue

            was_created = self.store.upsert_summary(
                PeriodType.HOURLY,
                key={"date": summary_date, "hour": hr},
                summary=summary_text,
                message_count=total_messages,
            )
            if was_created:
                created += 1

        return created

    # ------------------------------------------------------------------
    # Daily summaries
    # ------------------------------------------------------------------

    async def _update_daily(self, current_time: datetime) -> int:
        tz_now = force_timezone(current_time, self.timezone)
        today = tz_now.date()

        existing_keys = self.store.get_summary_keys(PeriodType.DAILY)
        existing_dates = {k[0] for k in existing_keys}

        hourly_summaries = self.store.get_summaries_before(PeriodType.HOURLY, today)

        hourly_by_day: dict[date, list[SummaryRecord]] = {}
        for hs in hourly_summaries:
            hourly_by_day.setdefault(hs.date, []).append(hs)

        pending = []
        for day, hour_summaries in hourly_by_day.items():
            if day not in existing_dates:
                parts = [f"[{(hs.hour or 0):02d}:00] {hs.summary}" for hs in hour_summaries]
                combined = "\n\n".join(parts)
                total = sum(hs.message_count for hs in hour_summaries)
                pending.append((day, combined, total))

        if not pending:
            return 0

        existing_memory = await self._get_context_for_summary(
            force_timezone(current_time, self.timezone), PeriodType.DAILY
        )

        created = 0
        for day, content, total_messages in sorted(pending):
            content = self._trim_content(content)
            summary_text = await self._create_summary(
                content=content,
                period_type=PeriodType.DAILY,
                existing_memory=existing_memory,
            )
            if not summary_text or not summary_text.strip():
                continue

            was_created = self.store.upsert_summary(
                PeriodType.DAILY,
                key={"date": day},
                summary=summary_text,
                message_count=total_messages,
            )
            if was_created:
                created += 1

        return created

    # ------------------------------------------------------------------
    # Weekly summaries
    # ------------------------------------------------------------------

    async def _update_weekly(self, current_time: datetime) -> int:
        tz_now = force_timezone(current_time, self.timezone)
        today = tz_now.date()
        this_week_start = week_start_sunday(datetime.combine(today, time(0, 0, 0))).date()

        dailies = self.store.get_summaries_before(PeriodType.DAILY, this_week_start)

        weekly_groups: dict[date, list[SummaryRecord]] = {}
        for d in dailies:
            week = week_start_sunday(datetime.combine(d.date, time(0, 0, 0))).date()
            weekly_groups.setdefault(week, []).append(d)

        existing_keys = self.store.get_summary_keys(PeriodType.WEEKLY)
        existing_weeks = {k[0] for k in existing_keys}

        pending = []
        for week, dailies_in_week in weekly_groups.items():
            if week not in existing_weeks and dailies_in_week:
                parts = [f"[{d.date}] {d.summary}" for d in dailies_in_week]
                combined = "\n\n".join(parts)
                total = sum(d.message_count for d in dailies_in_week)
                pending.append((week, combined, total))

        if not pending:
            return 0

        existing_memory = await self._get_context_for_summary(
            force_timezone(current_time, self.timezone), PeriodType.WEEKLY
        )

        created = 0
        for week, content, total_messages in sorted(pending):
            content = self._trim_content(content)
            summary_text = await self._create_summary(
                content=content,
                period_type=PeriodType.WEEKLY,
                existing_memory=existing_memory,
            )
            if not summary_text or not summary_text.strip():
                continue

            was_created = self.store.upsert_summary(
                PeriodType.WEEKLY,
                key={"week_start_date": week},
                summary=summary_text,
                message_count=total_messages,
            )
            if was_created:
                created += 1

        return created

    # ------------------------------------------------------------------
    # Monthly summaries
    # ------------------------------------------------------------------

    async def _update_monthly(self, current_time: datetime) -> int:
        tz_now = force_timezone(current_time, self.timezone)
        this_month_start = tz_now.date().replace(day=1)

        weeklies = self.store.get_summaries_before(PeriodType.WEEKLY, this_month_start)

        monthly_groups: dict[tuple[int, int], list[SummaryRecord]] = {}
        for w in weeklies:
            wsd = w.week_start_date
            if wsd:
                key = (wsd.year, wsd.month)
                monthly_groups.setdefault(key, []).append(w)

        existing_keys = self.store.get_summary_keys(PeriodType.MONTHLY)

        pending = []
        for (year, month), weeklies_in_month in monthly_groups.items():
            if (year, month) not in existing_keys and weeklies_in_month:
                parts = [f"[Week of {w.week_start_date}] {w.summary}" for w in weeklies_in_month]
                combined = "\n\n".join(parts)
                total = sum(w.message_count for w in weeklies_in_month)
                pending.append(((year, month), combined, total))

        if not pending:
            return 0

        existing_memory = await self._get_context_for_summary(
            force_timezone(current_time, self.timezone), PeriodType.MONTHLY
        )

        created = 0
        for (year, month), content, total_messages in sorted(pending):
            content = self._trim_content(content)
            summary_text = await self._create_summary(
                content=content,
                period_type=PeriodType.MONTHLY,
                existing_memory=existing_memory,
            )
            if not summary_text or not summary_text.strip():
                continue

            was_created = self.store.upsert_summary(
                PeriodType.MONTHLY,
                key={"year": year, "month": month},
                summary=summary_text,
                message_count=total_messages,
            )
            if was_created:
                created += 1

        return created

    # ------------------------------------------------------------------
    # Shared helpers
    # ------------------------------------------------------------------

    async def _get_context_for_summary(
        self,
        reference_time: datetime,
        min_resolution: PeriodType,
    ) -> str | None:
        builder = ContextBuilder(
            store=self.store,
            formatter=self.formatter,
            timezone=self.timezone,
        )
        data = builder.fetch_data(reference_time, min_resolution)
        context = builder.render(data, max_tokens=10000)
        return context if context and context.strip() else None

    async def _create_summary(
        self,
        *,
        content: str,
        period_type: PeriodType,
        existing_memory: str | None = None,
        max_retries: int = 3,
    ) -> str | None:
        meta = PERIOD_META[period_type]

        prompt = (
            f"You are an AI assistant creating a memory summary from {meta.time_period}. "
            'Write in FIRST PERSON ("I worked with...", "I helped with...").\n\n'
            f"Keep your summary to {meta.max_length} maximum.\n\n"
        )

        if existing_memory:
            prompt += (
                "Here is your existing memory context:\n"
                "-- BEGIN EXISTING MEMORY --\n"
                f"{existing_memory}\n"
                "-- END EXISTING MEMORY --\n\n"
            )

        footer = (
            "For your summary focus on:\n"
            f"{meta.focus}\n\n"
            "CRITICAL: Only use information explicitly stated above to summarize or infer "
            "learnings. Do NOT add details.\n\n"
            "This is the new content that you must summarize:\n"
        )

        enc = tiktoken.get_encoding("o200k_base")
        max_total_tokens = 200_000
        fixed_parts = prompt + footer
        fixed_tokens = len(enc.encode(fixed_parts))
        available = max_total_tokens - fixed_tokens - 10
        content_tokens = len(enc.encode(content))

        if content_tokens > available:
            content = enc.decode(enc.encode(content)[:available])

        prompt += footer + f"{content}\n\n"

        last_err: Exception | None = None
        for attempt in range(max_retries):
            try:
                resp: LLMResponse = await self.llm.complete(prompt, self.model, max_tokens=4096)
                self.total_usage["input_tokens"] += resp.input_tokens
                self.total_usage["output_tokens"] += resp.output_tokens
                self.total_usage["llm_calls"] += 1
                return resp.text.strip() if resp.text else None
            except Exception as exc:
                last_err = exc
                await asyncio.sleep(min(2**attempt, 8) + (attempt * 0.1))

        logger.error("Summary creation failed after %d attempts: %s", max_retries, last_err)
        return None

    def _prepare_content(self, messages_data: list[dict], max_tokens: int = 185_000) -> str:
        if not messages_data:
            return json.dumps([])

        enc = tiktoken.get_encoding("o200k_base")
        full = json.dumps(messages_data)
        if len(enc.encode(full)) <= max_tokens:
            return full

        selected = []
        current_tokens = 2
        for msg in messages_data:
            msg_json = json.dumps(msg)
            tokens = len(enc.encode(msg_json))
            if current_tokens + tokens + 2 <= max_tokens:
                selected.append(msg)
                current_tokens += tokens + 2
            else:
                break

        return json.dumps(selected)

    def _trim_content(self, content: str, max_tokens: int = 185_000) -> str:
        enc = tiktoken.get_encoding("o200k_base")
        tokens = enc.encode(content)
        if len(tokens) <= max_tokens:
            return content
        return enc.decode(tokens[:max_tokens])
