"""Summarizer — rolling conversation summary with turn-based triggers."""
from __future__ import annotations

import logging
from typing import Any, Callable, Awaitable

from kokoro_chat.config import ChatConfig, calculate_cost
from kokoro_chat.storage.protocol import (
    CacheStore,
    ChatHistoryStore,
    RollingSummaryStore,
    UsageLogStore,
)

logger = logging.getLogger(__name__)

AiChatFn = Callable[..., Awaitable[Any]]


class Summarizer:
    def __init__(
        self,
        summary_store: RollingSummaryStore,
        history_store: ChatHistoryStore,
        cache: CacheStore,
        ai_chat_fn: AiChatFn,
        config: ChatConfig,
        usage_log: UsageLogStore | None = None,
    ) -> None:
        self._summary_store = summary_store
        self._history = history_store
        self._cache = cache
        self._ai_chat = ai_chat_fn
        self._config = config
        self._usage_log = usage_log

    async def maybe_summarize(self, char_id: str, user_id: str, character_name: str) -> None:
        """Increment turn count. Summarize if threshold reached."""
        cached_summary = await self._cache.get_rolling_summary(user_id, char_id)
        if cached_summary is not None:
            turn_count = cached_summary.get("turn_count", 0)
            current_summary = cached_summary.get("summary", "")
        else:
            record = await self._summary_store.get_or_create(user_id, char_id)
            turn_count = record.turn_count
            current_summary = record.summary or ""

        turn_count += 1

        if turn_count > 0 and turn_count % self._config.summarize_every == 0:
            current_summary = await self._do_summarize(
                char_id, user_id, character_name,
                current_summary, turn_count,
            )

        await self._summary_store.update(
            user_id, char_id, current_summary, turn_count,
        )
        await self._cache.set_rolling_summary(
            user_id, char_id, current_summary, turn_count,
        )

    async def _do_summarize(
        self,
        char_id: str,
        user_id: str,
        character_name: str,
        existing_summary: str,
        turn_count: int,
    ) -> str:
        """Compress older messages into rolling summary via LLM."""
        cfg = self._config
        limit = 40
        chat_msgs = await self._history.get_recent(user_id, char_id, limit=limit)
        if len(chat_msgs) < 10:
            return existing_summary

        to_compress = chat_msgs[:-10]
        dialogue_text = "\n".join(
            f"{'用戶' if m.role == 'user' else character_name}：{m.content}"
            for m in to_compress[-20:]
        )
        prompt = f"""請將以下對話歷史壓縮成一段簡潔的摘要（200字以內）。
保留重要的：事件、承諾、情感變化、關係進展、用戶提到的個人資訊。

{"之前的摘要：" + existing_summary if existing_summary else ""}

需要壓縮的對話：
{dialogue_text}

只輸出摘要，不要其他內容。"""

        try:
            response = await self._ai_chat(
                [{"role": "user", "content": prompt}],
                model=cfg.memory_llm_model,
                max_tokens=300,
            )
            new_summary = response.content.strip()

            if self._usage_log is not None:
                try:
                    cost = calculate_cost(
                        cfg.memory_llm_model,
                        response.usage.input_tokens,
                        response.usage.output_tokens,
                    )
                    await self._usage_log.log(
                        user_id=user_id,
                        model=cfg.memory_llm_model,
                        input_tokens=response.usage.input_tokens,
                        output_tokens=response.usage.output_tokens,
                        cost_usd=cost,
                        task_type="summarize",
                        character_id=char_id,
                    )
                except Exception as e:
                    logger.warning("Failed to log usage: %s", e)

            await self._summary_store.update(
                user_id, char_id, new_summary, turn_count,
            )
            await self._cache.set_rolling_summary(
                user_id, char_id, new_summary, turn_count,
            )

            recent_msgs = [
                {"role": m.role, "content": m.content}
                for m in chat_msgs[-10:]
            ]
            await self._cache.set_recent_messages(user_id, char_id, recent_msgs)

            logger.info("Rolling summary updated (%d messages compressed)", len(to_compress))
            return new_summary
        except Exception as e:
            logger.warning("Summarization failed: %s", e, exc_info=True)
            return existing_summary
