"""ChatEngine core — extracted from app/services/chat_engine.py.

Accepts storage via dependency injection (Protocol-based stores).
No sqlalchemy / fastapi / redis imports allowed in this file.

The engine owns:
  - Attitude state management (load, save, decay)
  - Conversation history retrieval and persistence
  - System prompt assembly orchestration
  - Background task scheduling (summarize, decay)
  - LLM client construction

The app layer is responsible for:
  - Building concrete storage implementations
  - Providing ChatConfig with credentials and settings
  - Wiring the engine into HTTP/WebSocket routes
"""
from __future__ import annotations

import asyncio
import logging
import time
from typing import Any

from motosan_ai import Client as AiClient
from motosan_chat import (
    AgentLoop,
    IncomingEvent,
    Message as McMessage,
    Thread as McThread,
)

from kokoro_chat.attitude.state import AttitudeState
from kokoro_chat.character.model import Character
from kokoro_chat.config import ChatConfig, calculate_cost
from kokoro_chat.engine.background import strip_app_markup
from kokoro_chat.engine.attitude_manager import AttitudeManager
from kokoro_chat.engine.message_manager import MessageManager
from kokoro_chat.engine.summarizer import Summarizer
from kokoro_chat.engine.llm import (
    STREAM_SENTINEL,
    StreamingChannel,
    UsageTrackingAiClient,
    StreamingMotosanAiClient,
    build_ai_client,
)
from kokoro_chat.prompt.assembler import (
    assemble_static_system_prompt,
    assemble_static_system_prompt_blocks,
    assemble_system_prompt,
)
from kokoro_chat.storage.protocol import (
    AttitudeStore,
    CacheStore,
    ChatHistoryStore,
    ChatMessageWriter,
    EventStore,
    MemoStore,
    MemoryVectorStore,
    RollingSummaryStore,
    UsageLogStore,
)

logger = logging.getLogger(__name__)


# ══════════════════════════════════════════════════════════════════════════════
#  ChatEngine — core class with DI storage
# ══════════════════════════════════════════════════════════════════════════════

class ChatEngine:
    """Core chat engine with dependency-injected storage.

    The engine does NOT know about databases, Redis, or web frameworks.
    All persistence is mediated through Protocol-based stores injected
    at construction time.

    Parameters
    ----------
    config : ChatConfig
        Pure-data configuration (models, tokens, feature flags, credentials).
    attitude_store : AttitudeStore
        Read/write attitude state.
    cache : CacheStore
        Fast cache layer (attitude, messages, summaries).
    chat_history : ChatHistoryStore
        Read recent chat messages.
    chat_writer : ChatMessageWriter
        Write new chat messages.
    memory_store : MemoryVectorStore
        Semantic search and persistence for memories.
    memo_store : MemoStore
        Key-value memo storage.
    rolling_summary_store : RollingSummaryStore
        Read/write rolling conversation summaries.
    event_store : EventStore
        Read events for a character/user pair.
    usage_log : UsageLogStore | None
        Optional usage logging.
    create_tools_fn : callable | None
        Optional factory for creating motosan-chat Tool instances.
        Signature: (char_id: str) -> list[Tool].
        If not provided, the engine will have no tool support.
    story_engine : Any | None
        Optional story engine for story-mode context.
    """

    def __init__(
        self,
        config: ChatConfig,
        *,
        attitude_store: AttitudeStore,
        cache: CacheStore,
        chat_history: ChatHistoryStore,
        chat_writer: ChatMessageWriter,
        memory_store: MemoryVectorStore,
        memo_store: MemoStore,
        rolling_summary_store: RollingSummaryStore,
        event_store: EventStore,
        usage_log: UsageLogStore | None = None,
        create_tools_fn: Any = None,
        story_engine: Any = None,
    ):
        self.config = config
        self._attitude_store = attitude_store
        self._cache = cache
        self._chat_history = chat_history
        self._chat_writer = chat_writer
        self._memory_store = memory_store
        self._memo_store = memo_store
        self._rolling_summary_store = rolling_summary_store
        self._event_store = event_store
        self._usage_log = usage_log
        self._create_tools_fn = create_tools_fn
        self.story_engine = story_engine

        # Build LLM clients
        self.chat_llm = build_ai_client(
            config.chat_provider, config.chat_model, config,
        )
        self.memory_llm = build_ai_client(
            config.memory_provider, config.memory_model, config,
        )
        if self.chat_llm is None:
            self.memory_llm = None
        elif self.memory_llm is None:
            logger.warning(
                "Memory provider unavailable (%s); fallback to chat provider (%s)",
                config.memory_provider, config.chat_provider,
            )
            self.memory_llm = self.chat_llm

        self._background_tasks: set[asyncio.Task] = set()
        self._bg_semaphore = asyncio.Semaphore(config.bg_task_concurrency)

        # Attitude manager
        self._attitude_mgr = AttitudeManager(attitude_store, cache)

        # Message manager
        self._msg_mgr = MessageManager(
            chat_history, chat_writer, cache, rolling_summary_store, config,
        )

        # Summarizer
        self._summarizer = Summarizer(
            rolling_summary_store, chat_history, cache,
            self._ai_chat, config, usage_log,
        )

    # -- LLM client accessors --------------------------------------------------

    def _get_chat_llm_client(self) -> AiClient:
        if self.chat_llm:
            return self.chat_llm
        raise RuntimeError("No chat model provider credentials available")

    def _get_memory_llm_client(self) -> AiClient:
        if self.memory_llm:
            return self.memory_llm
        return self._get_chat_llm_client()

    async def _ai_chat(
        self,
        messages: list[dict[str, str]],
        *,
        model: str,
        max_tokens: int,
        system: str | None = None,
        client_kind: str = "chat",
    ) -> Any:
        client = (
            self._get_memory_llm_client()
            if client_kind == "memory"
            else self._get_chat_llm_client()
        )
        return await client.chat(
            messages,
            system=system,
            max_tokens=max_tokens,
            provider_options={"model": model},
        )

    # -- Usage logging ---------------------------------------------------------

    async def _log_usage(
        self,
        user_id: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        task_type: str,
        character_id: str | None = None,
    ) -> None:
        if self._usage_log is None:
            return
        try:
            cost = calculate_cost(model, input_tokens, output_tokens)
            await self._usage_log.log(
                user_id=user_id,
                model=model,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cost_usd=cost,
                task_type=task_type,
                character_id=character_id,
            )
        except Exception as e:
            logger.warning("Failed to log usage: %s", e)

    # -- Background tasks ------------------------------------------------------

    async def _run_background_tasks(
        self,
        user_message: str,
        assistant_reply: str,
        character: Character,
        user_id: str,
        char_id: str,
        attitude: AttitudeState,
    ) -> None:
        async with self._bg_semaphore:
            try:
                fresh_attitude = await self._attitude_mgr.load(char_id, user_id)
                fresh_attitude.decay_modifiers()
                await self._attitude_mgr.save(fresh_attitude)
                await self._summarizer.maybe_summarize(char_id, user_id, character.name)
            except Exception as e:
                logger.warning("Background tasks failed: %s", e, exc_info=True)

    # -- Context preparation ---------------------------------------------------

    async def _prepare_slim_chat_context(
        self,
        user_message: str,
        character: Character,
        user_id: str,
        chat_mode: str,
        scene: str | None,
    ) -> dict:
        """Prepare slim context for AgentLoop path."""
        char_id = str(character.id)
        cfg = self.config

        attitude = await self._attitude_mgr.load(char_id, user_id)
        _rolling_summary, recent_messages = await self._msg_mgr.get_for_prompt(
            char_id, user_id,
        )

        system_prompt = assemble_static_system_prompt(
            character=character,
            chat_mode=chat_mode,
            scene=scene,
        )
        if attitude.cooldown_hint_needed:
            system_prompt += (
                "\n\n[關係冷卻提示]\n"
                "你們已經有一段時間沒有聯繫，角色心情可能略顯疏遠，"
                "但仍保留過往關係基礎。"
            )

        system_prompt_blocks: list[dict] | None = None
        if cfg.prompt_cache_enabled:
            system_prompt_blocks = assemble_static_system_prompt_blocks(
                character=character,
                chat_mode=chat_mode,
                scene=scene,
            )
            if attitude.cooldown_hint_needed:
                system_prompt_blocks.append({
                    "type": "text",
                    "text": (
                        "\n\n[關係冷卻提示]\n"
                        "你們已經有一段時間沒有聯繫，角色心情可能略顯疏遠，"
                        "但仍保留過往關係基礎。"
                    ),
                })

        recent_messages = MessageManager.build_few_shot_prefix(character) + recent_messages

        await self._msg_mgr.add_message(char_id, user_id, "user", user_message, chat_mode)
        recent_messages.append({"role": "user", "content": user_message})

        if character.post_history_instructions:
            recent_messages.append({
                "role": "user",
                "content": f"[系統提醒：{character.post_history_instructions}]",
            })

        return {
            "char_id": char_id,
            "attitude": attitude,
            "system_prompt": system_prompt,
            "system_prompt_blocks": system_prompt_blocks,
            "recent_messages": recent_messages,
        }

    async def _prepare_chat_context(
        self,
        user_message: str,
        character: Character,
        user_id: str,
        chat_mode: str,
        scene: str | None,
        temporary_event: str | None,
    ) -> dict:
        """Prepare full context with all providers (for non-AgentLoop path)."""
        char_id = str(character.id)

        # Parallel: DB queries + memory search
        async def _store_queries():
            attitude = await self._attitude_mgr.load(char_id, user_id)
            events_raw = await self._event_store.get_by_character(char_id, user_id)
            memos_raw = await self._memo_store.get(user_id, char_id)
            rolling_summary, recent_messages = await self._msg_mgr.get_for_prompt(
                char_id, user_id,
            )
            return attitude, events_raw, memos_raw, rolling_summary, recent_messages

        store_result, memory_result = await asyncio.gather(
            _store_queries(),
            self._memory_store.search(
                query=user_message, character_id=char_id, user_id=user_id,
            ),
            return_exceptions=True,
        )

        if isinstance(store_result, BaseException):
            raise store_result
        if isinstance(memory_result, BaseException):
            logger.warning(
                "Memory provider failed, continue without memories: %s",
                memory_result,
            )
            memory_result = []

        attitude, events_raw, memos_raw, rolling_summary, recent_messages = store_result
        memories = memory_result

        events = [{"type": e.event_type, "content": e.content} for e in events_raw]
        if temporary_event:
            events.append({"type": "temporary", "content": temporary_event})
        memos = [
            {"key": m.get("key", ""), "value": m.get("value", "")}
            for m in memos_raw
        ] if isinstance(memos_raw, list) else []

        system_prompt = assemble_system_prompt(
            character=character,
            affection_score=attitude.affection,
            chat_mode=chat_mode,
            events=events,
            memories=[
                {"text": m.text, "relevance": m.relevance}
                for m in memories
            ] if memories else [],
            memos=memos,
            scene=scene,
            attitude_section=attitude.to_prompt_section(),
            rolling_summary=rolling_summary,
            trust_score=attitude.trust,
        )
        if attitude.cooldown_hint_needed:
            system_prompt += (
                "\n\n[關係冷卻提示]\n"
                "你們已經有一段時間沒有聯繫，角色心情可能略顯疏遠，"
                "但仍保留過往關係基礎。"
            )

        recent_messages = MessageManager.build_few_shot_prefix(character) + recent_messages

        await self._msg_mgr.add_message(char_id, user_id, "user", user_message, chat_mode)
        recent_messages.append({"role": "user", "content": user_message})

        if character.post_history_instructions:
            recent_messages.append({
                "role": "user",
                "content": f"[系統提醒：{character.post_history_instructions}]",
            })

        return {
            "char_id": char_id,
            "attitude": attitude,
            "system_prompt": system_prompt,
            "recent_messages": recent_messages,
        }

    # -- Public API: chat (non-streaming) --------------------------------------

    async def chat(
        self,
        user_message: str,
        character: Character,
        user_id: str,
        chat_mode: str = "short",
        scene: str | None = None,
        temporary_event: str | None = None,
    ) -> dict:
        """Complete chat flow -- collects all streamed events into a single dict.

        Returns: {"reply": str, "attitude": dict}
        """
        reply_chunks: list[str] = []
        attitude = None
        async for event in self.chat_stream(
            user_message, character, user_id,
            chat_mode=chat_mode, scene=scene,
            temporary_event=temporary_event,
        ):
            if event["type"] == "text_delta":
                reply_chunks.append(event["text"])
            elif event["type"] == "message_complete":
                attitude = event.get("attitude")
        return {"reply": "".join(reply_chunks), "attitude": attitude}

    # -- Public API: chat_stream (streaming) -----------------------------------

    async def chat_stream(
        self,
        user_message: str,
        character: Character,
        user_id: str,
        chat_mode: str = "short",
        scene: str | None = None,
        temporary_event: str | None = None,
    ):
        """Streaming chat flow (async generator).

        yields: {"type": "text_delta", "text": "..."} or
                {"type": "message_complete", "character_name": ..., "attitude": {...}}
        """
        cfg = self.config
        _t0 = time.monotonic()

        ctx = await self._prepare_slim_chat_context(
            user_message, character, user_id, chat_mode, scene,
        )
        char_id = ctx["char_id"]
        attitude = ctx["attitude"]

        _t1 = time.monotonic()
        logger.info("[timing] prepare_context: %.1fms", (_t1 - _t0) * 1000)

        # Build tools
        tools = []
        if self._create_tools_fn is not None:
            tools = self._create_tools_fn(char_id)

        tracking_ai = UsageTrackingAiClient(
            self._get_chat_llm_client(),
            model=cfg.chat_model,
            max_tokens=cfg.max_tokens,
        )
        llm_client = StreamingMotosanAiClient(
            tracking_ai,
            system=ctx["system_prompt"],
            system_prompt_blocks=ctx.get("system_prompt_blocks"),
            prompt_cache_enabled=cfg.prompt_cache_enabled,
        )

        mc_messages = [
            McMessage.user(m["content"])
            if m["role"] == "user"
            else McMessage.assistant(m["content"])
            for m in ctx["recent_messages"]
        ]

        chunk_queue: asyncio.Queue = asyncio.Queue()
        streaming_channel = StreamingChannel(chunk_queue)
        fake_event = IncomingEvent(
            user_id=user_id,
            content="",
            channel_name="api",
        )
        thread = McThread(event=fake_event, channel=streaming_channel)
        agent_loop = AgentLoop(tools=tools, max_iterations=5)

        agent_result_holder: list = []
        agent_error_holder: list = []

        async def _run_agent():
            try:
                result = await agent_loop.run(llm_client, thread, mc_messages)
                agent_result_holder.append(result)
            except RuntimeError as exc:
                logger.warning("chat_stream AgentLoop max iterations: %s", exc)
                agent_error_holder.append(exc)
            except Exception as exc:
                logger.error("chat_stream AgentLoop error: %s", exc, exc_info=True)
                agent_error_holder.append(exc)
            finally:
                await chunk_queue.put(STREAM_SENTINEL)

        _t2 = time.monotonic()
        logger.info("[timing] build_agent: %.1fms", (_t2 - _t1) * 1000)

        agent_task = asyncio.create_task(_run_agent())

        full_text_parts: list[str] = []
        _t_first_chunk = None
        while True:
            item = await chunk_queue.get()
            if item is STREAM_SENTINEL:
                break
            if _t_first_chunk is None:
                _t_first_chunk = time.monotonic()
                logger.info(
                    "[timing] first_chunk (TTFB): %.1fms",
                    (_t_first_chunk - _t2) * 1000,
                )
            full_text_parts.append(item)
            yield {"type": "text_delta", "text": item}

        await agent_task

        _t3 = time.monotonic()
        logger.info("[timing] llm_streaming: %.1fms", (_t3 - _t2) * 1000)

        full_text = strip_app_markup("".join(full_text_parts))

        # Log tool calls
        if agent_result_holder:
            result = agent_result_holder[0]
            for tool_name, tool_args in result.tool_calls:
                logger.info("[stream] Tool called: %s(%s)", tool_name, tool_args)

        # Empty reply fallback
        if not full_text:
            logger.info("Empty reply after tool calls -- requesting text-only followup")
            followup_messages = list(mc_messages)
            followup_messages.append(
                McMessage.user("[system: tools executed, now reply to the user in character]"),
            )
            try:
                followup_resp = await llm_client.chat(followup_messages, [])
                full_text = getattr(followup_resp, "content", "") or ""
                if full_text:
                    yield {"type": "text_delta", "text": full_text}
            except Exception as exc:
                logger.warning("Fallback reply failed: %s", exc)

        # Log usage
        await self._log_usage(
            user_id, cfg.chat_model,
            tracking_ai.total_input_tokens,
            tracking_ai.total_output_tokens,
            "chat_stream", character_id=char_id,
        )

        # Save assistant reply
        await self._msg_mgr.add_message(char_id, user_id, "assistant", full_text, chat_mode)

        _t4 = time.monotonic()
        logger.info(
            "[timing] TOTAL: %.1fms (context=%.0f, llm=%.0f, save=%.0f)",
            (_t4 - _t0) * 1000, (_t1 - _t0) * 1000,
            (_t3 - _t2) * 1000, (_t4 - _t3) * 1000,
        )

        # Background tasks
        attitude_snapshot = attitude.snapshot()
        bg_coro = self._run_background_tasks(
            user_message, full_text, character, user_id, char_id, attitude,
        )
        task = asyncio.create_task(bg_coro)
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)

        yield {
            "type": "message_complete",
            "character_name": character.name,
            "attitude": attitude_snapshot,
        }

    # -- Public API: get_attitude ----------------------------------------------

    async def get_attitude(self, char_id: str, user_id: str) -> dict:
        """Query current attitude state."""
        attitude = await self._attitude_mgr.load(char_id, user_id)
        return attitude.snapshot()

    # -- Public API: trigger_event ---------------------------------------------

    async def trigger_event(
        self,
        character: Character,
        user_id: str,
        event_description: str,
    ) -> dict:
        """Manually trigger an event that changes character attitude."""
        from kokoro_chat.engine.background import (
            build_event_analysis_prompt,
            parse_event_impact,
        )

        char_id = str(character.id)
        cfg = self.config
        attitude = await self._attitude_mgr.load(char_id, user_id)

        prompt = build_event_analysis_prompt(
            character_name=character.name,
            personality=character.personality,
            affection=attitude.affection,
            trust=attitude.trust,
            mood=attitude.mood,
            event_description=event_description,
        )

        try:
            response = await self._ai_chat(
                [{"role": "user", "content": prompt}],
                model=cfg.memory_llm_model,
                max_tokens=200,
            )
            await self._log_usage(
                user_id, cfg.memory_llm_model,
                response.usage.input_tokens, response.usage.output_tokens,
                "event_analysis", character_id=char_id,
            )

            raw_text = response.content.strip() if response.content else ""
            impact = parse_event_impact(raw_text)
            if impact is None:
                logger.warning("Event analysis: cannot parse impact from: %s", raw_text[:200])
                return attitude.snapshot()

            attitude.apply_event_impact(event_description, impact)

            await self._memory_store.add(
                character_id=char_id,
                user_id=user_id,
                text=f"重大事件：{event_description}。態度變化：{impact.get('attitude_shift', '')}",
                memory_type="attitude_shift",
                importance=0.95,
            )
            await self._attitude_mgr.save(attitude)

            logger.info("Event triggered attitude shift: %s", impact)
            return attitude.snapshot()
        except Exception as e:
            logger.warning("Event analysis failed: %s", e, exc_info=True)
            return attitude.snapshot()
