"""ProcessMessage workflow — main message processing."""

from __future__ import annotations

from contextlib import AsyncExitStack
from typing import Any, TYPE_CHECKING

from loguru import logger

if TYPE_CHECKING:
    from miu_bot.db.backend import MemoryBackend


def _annotate_sender(content: str, sender_id: str, sender_name: str) -> str:
    """Prefix message with sender info so LLM can identify who is speaking."""
    parts = []
    if sender_id:
        parts.append(f"userId={sender_id}")
    if sender_name:
        parts.append(sender_name)
    if not parts:
        return content
    return f"[{' / '.join(parts)}]: {content}"


def _build_history(messages: list, is_group: bool) -> list[dict[str, Any]]:
    """Build LLM history, annotating user messages with sender info in groups."""
    history = []
    for m in messages:
        if m.role == "user" and is_group and m.metadata:
            meta = m.metadata if isinstance(m.metadata, dict) else {}
            sid = meta.get("sender_id", "")
            sname = meta.get("sender_name", "")
            if sid or sname:
                history.append({"role": m.role, "content": _annotate_sender(m.content, sid, sname)})
                continue
        history.append({"role": m.role, "content": m.content})
    return history


class ProcessMessageWorkflow:
    """Workflow for processing a single inbound message.

    Per-message: creates provider, tools, MCP from workspace.config_overrides.
    """

    def __init__(
        self,
        backend: "MemoryBackend",
        gateway_url: str = "http://localhost:18790",
        # Fallback defaults (used when workspace has no provider override)
        fallback_model: str = "anthropic/claude-opus-4-6",
        fallback_api_key: str = "",
        fallback_api_base: str | None = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        max_iterations: int = 20,
    ):
        self.backend = backend
        self.gateway_url = gateway_url
        self.fallback_model = fallback_model
        self.fallback_api_key = fallback_api_key
        self.fallback_api_base = fallback_api_base
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.max_iterations = max_iterations

    async def process(self, workflow_input: dict[str, Any]) -> dict[str, Any]:
        """Process a message:received event."""
        import time as _time

        from temporalio import activity

        from miu_bot.observability.spans import get_tracer
        from miu_bot.workspace.identity import parse_identity, render_system_prompt
        from miu_bot.agent.context import ContextBuilder
        from miu_bot.agent.processor import run_agent_loop
        from miu_bot.agent.tools.registry import ToolRegistry
        from miu_bot.worker.workflows.shared import (
            create_provider, connect_mcp, connect_skill_mcp, sanitize_response,
        )
        t_start = _time.monotonic()

        # Send heartbeat during long setup phases (MCP, context assembly)
        def _heartbeat(data: dict) -> None:
            try:
                activity.heartbeat(data)
            except Exception:
                pass  # heartbeat failures are non-fatal

        workspace_id = workflow_input["workspace_id"]
        channel = workflow_input["channel"]
        chat_id = workflow_input["chat_id"]
        content = workflow_input["content"]
        metadata = workflow_input.get("metadata", {})
        sender_id = workflow_input.get("sender_id", "")
        bot_name = workflow_input.get("bot_name", "")
        media_urls = workflow_input.get("media", [])

        # Persist sender_id in metadata for history attribution
        if sender_id:
            metadata["sender_id"] = sender_id

        tracer = get_tracer()

        # Load workspace
        workspace = await self.backend.get_workspace(workspace_id)
        if not workspace or workspace.status != "active":
            return {"status": "skipped", "reason": "workspace_inactive"}

        span_ctx = tracer.start_as_current_span(
            "process_message",
            attributes={
                "miubot.bot": bot_name,
                "miubot.channel": channel,
                "miubot.workspace_id": workspace_id,
            },
        ) if tracer else None

        span = span_ctx.__enter__() if span_ctx else None

        # Ensure session exists (keyed by workspace + channel + chat_id)
        session = await self.backend.get_or_create_session(
            workspace_id, channel, chat_id,
        )
        session_id = session.id

        # Create per-workspace provider (shared helper)
        provider, model = create_provider(
            workspace.config_overrides,
            self.fallback_model, self.fallback_api_key, self.fallback_api_base,
        )
        if span:
            span.set_attribute("miubot.model", model)

        # Create per-workspace tools (MCP)
        tools = ToolRegistry()
        mcp_stack = AsyncExitStack()
        await mcp_stack.__aenter__()
        try:
            # Load templates and skills from DB
            templates = await self.backend.get_templates(workspace_id)
            db_skills = await self.backend.get_skills(workspace_id)

            # Connect MCP servers from config_overrides (shared helper)
            mcp_count = await connect_mcp(
                workspace.config_overrides, tools, mcp_stack
            )
            if mcp_count:
                logger.info(f"Connected {mcp_count} MCP server(s) for {bot_name}")
                if span:
                    span.set_attribute("miubot.mcp_servers", mcp_count)
            _heartbeat({"phase": "mcp_connected", "mcp_count": mcp_count})

            # Load context
            messages = await self.backend.get_messages(session_id, limit=50)

            # Use tier-based context assembly (BASB)
            from miu_bot.memory.context_assembly import assemble_memory_context
            from miu_bot.workspace.identity import compose_from_templates

            memories_text = await assemble_memory_context(
                self.backend, workspace_id, query=content
            )
            history = _build_history(messages, is_group=metadata.get("is_group", False))

            # Compose base prompt: templates (new) or identity (legacy)
            if templates:
                base_prompt = compose_from_templates(templates, memories_text)
            else:
                identity = parse_identity(workspace.identity)
                base_prompt = render_system_prompt(identity, memories_text)

            # Merge skills: DB (new) or config_overrides (legacy fallback)
            skills_section = ""
            skill_mcp: dict = {}
            if db_skills:
                from miu_bot.skills.merger import merge_skills_from_db
                skills_section, skill_mcp, _ = merge_skills_from_db(db_skills)
            else:
                skill_dicts = workspace.config_overrides.get("skills", [])
                if skill_dicts:
                    from miu_bot.skills.merger import merge_skills_into_prompt
                    from miu_bot.skills.schema import SkillConfig
                    skills = [SkillConfig.model_validate(s) for s in skill_dicts]
                    skills_section, skill_mcp, _ = merge_skills_into_prompt("", skills)

            # Combine base prompt + skills
            full_prompt = base_prompt
            if skills_section:
                full_prompt = f"{base_prompt}\n\n{skills_section}"

            # Connect skill-provided MCP servers (shared helper)
            if skill_mcp:
                mcp_count += await connect_skill_mcp(skill_mcp, tools, mcp_stack)

            # Register Zalo tool for Zalo channels (HTTP proxy to gateway)
            if channel == "zalo":
                from miu_bot.agent.tools.zalo import ZaloTool
                zalo_tool = ZaloTool(
                    gateway_url=self.gateway_url,
                    bot_name=bot_name,
                )
                zalo_tool.set_context(channel, chat_id)
                tools.register(zalo_tool)

            # Inject tool system hints into prompt
            tool_hints = tools.get_system_hints()
            if tool_hints:
                full_prompt = f"{full_prompt}\n\n{tool_hints}"

            # Annotate current message with sender for group chats
            current_message = content
            if metadata.get("is_group") and (sender_id or metadata.get("sender_name")):
                current_message = _annotate_sender(
                    content, sender_id, metadata.get("sender_name", ""),
                )

            # Resolve media URLs to base64/text for multimodal LLM input
            media_parts: list = []
            if media_urls:
                from miu_bot.agent.media_resolver import resolve_media_urls
                media_parts = await resolve_media_urls(media_urls)
                if media_parts:
                    logger.info(f"Resolved {len(media_parts)} media attachment(s) for {bot_name}")

            _heartbeat({"phase": "context_ready"})

            # Build LLM messages using composed prompt
            context_builder = ContextBuilder(workspace=None)
            llm_messages = context_builder.build_workspace_messages_from_prompt(
                prompt=full_prompt, history=history,
                current_message=current_message, channel=channel, chat_id=chat_id,
                is_group=metadata.get("is_group", False),
            )

            # Inject resolved media into the user message for multimodal processing
            if media_parts and llm_messages:
                last_msg = llm_messages[-1]
                if last_msg.get("role") == "user":
                    existing = last_msg["content"]
                    if isinstance(existing, str):
                        existing = [{"type": "text", "text": existing}]
                    llm_messages[-1]["content"] = media_parts + existing

            # Run agent loop with heartbeat reporting
            response_content, tools_used, trace = await run_agent_loop(
                provider=provider, messages=llm_messages, tools=tools,
                model=model, temperature=self.temperature,
                max_tokens=self.max_tokens,
                max_iterations=self.max_iterations,
                on_heartbeat=_heartbeat,
            )

            if response_content is None:
                response_content = (
                    "I've completed processing but have no response to give."
                )

            # Sanitize before sending to user (error boundary)
            response_content = sanitize_response(response_content)

            if span:
                span.set_attribute("miubot.tools_used", ",".join(tools_used))
                span.set_attribute("miubot.response_len", len(response_content))

            # Save messages
            await self.backend.save_message(
                session_id, "user", content, metadata
            )
            await self.backend.save_message(
                session_id, "assistant", response_content,
                {"tools_used": tools_used} if tools_used else None,
            )

            total_s = round(_time.monotonic() - t_start, 2)
            if span:
                span.set_attribute("miubot.total_s", total_s)

            # NOTE: Response is NOT sent here. The workflow sends it after
            # the activity succeeds, preventing duplicates on retries.
            return {
                "status": "ok",
                "bot": bot_name,
                "model": model,
                "tools_used": tools_used,
                "total_s": total_s,
                "response_content": response_content,
                "response_preview": response_content[:300],
                "trace": trace,
                # Pass through for workflow to send
                "channel": channel,
                "chat_id": chat_id,
                "metadata": metadata,
            }
        finally:
            if span_ctx:
                span_ctx.__exit__(None, None, None)
            try:
                await mcp_stack.aclose()
            except RuntimeError as e:
                if "cancel scope" in str(e):
                    logger.debug(f"Suppressed MCP cleanup error: {e}")
                else:
                    raise
