"""Conversation-brain contracts shared by every assistant type.

Brain selects who owns reasoning and conversation state.  The Pipecat pipeline
still owns media transport, STT/TTS, transcript delivery, and interruption
semantics.  This keeps assistant-specific orchestration out of pipeline.py
without coupling brains to Pipecat internals more than necessary.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Any, Protocol, runtime_checkable

from models import AssistantConfig
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.frames.frames import Frame
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameProcessor


@dataclass(frozen=True)
class BrainSpec:
    """Static capabilities used by validation and runtime dispatch."""

    type: str
    supported_runtime_modes: frozenset[str]
    # False means context, knowledge bases, and tools live on an external agent.
    owns_context: bool


class CallEndPort(Protocol):
    """Small call-lifecycle surface available to a brain."""

    @property
    def ending(self) -> bool: ...

    def begin(self, reason: str) -> None: ...

    def arm_after_speech(self) -> None: ...

    async def finish(self) -> None: ...


@dataclass(frozen=True)
class BrainRuntime:
    """Pipeline-owned capabilities injected into one brain session."""

    context: LLMContext
    llm: Any
    queue_frame: Callable[[Frame], Awaitable[None]]
    set_system_prompt: Callable[[str], None]
    set_tools: Callable[[list[FunctionSchema] | None], None]
    call_end: CallEndPort
    worker: Any = None
    context_aggregator: Any = None
    transport: Any = None
    switch_services: (
        Callable[[str | None, str | None, str | None], Awaitable[None]] | None
    ) = None
    set_knowledge_scope: Callable[[dict[str, Any]], None] | None = None
    set_input_enabled: Callable[[bool], None] | None = None
    apply_turn_config: (
        Callable[[bool, dict[str, Any]], Awaitable[None]] | None
    ) = None
    flow_global_functions: list[Any] = field(default_factory=list)


class BaseBrain:
    """No-op lifecycle defaults for brains without local orchestration."""

    spec: BrainSpec

    async def greeting(self, cfg: AssistantConfig) -> str:
        return cfg.greeting

    def system_prompt(self, cfg: AssistantConfig) -> str:
        return cfg.prompt if self.spec.owns_context else ""

    def build_llm(self, cfg: AssistantConfig, context: LLMContext) -> FrameProcessor:
        raise NotImplementedError

    async def setup(self, cfg: AssistantConfig, runtime: BrainRuntime) -> None:
        """Register tools and initialize per-call orchestration."""

    async def on_connected(self) -> None:
        """Handle a connected client after the common greeting is queued."""

    def record_user_message(self, content: str) -> None:
        """Observe a committed user message for brain-owned routing state."""

    async def on_user_turn_end(self, content: str) -> bool:
        """Handle a complete user turn before the conversational LLM runs.

        Return True when the brain scheduled the next action itself and the
        in-flight context frame must not reach the previous Agent's LLM.
        """
        self.record_user_message(content)
        return False

    async def on_assistant_text_start(self, turn_id: str) -> None:
        """Observe the start of a generated assistant turn."""

    async def on_assistant_text_end(
        self,
        turn_id: str,
        content: str,
        interrupted: bool,
    ) -> None:
        """Observe the completion of a generated assistant turn."""


@runtime_checkable
class Brain(Protocol):
    """One instance per call; implementations may keep conversation state."""

    spec: BrainSpec

    async def greeting(self, cfg: AssistantConfig) -> str: ...

    def system_prompt(self, cfg: AssistantConfig) -> str: ...

    def build_llm(self, cfg: AssistantConfig, context: LLMContext) -> FrameProcessor: ...

    async def setup(self, cfg: AssistantConfig, runtime: BrainRuntime) -> None: ...

    async def on_connected(self) -> None: ...

    def record_user_message(self, content: str) -> None: ...

    async def on_user_turn_end(self, content: str) -> bool: ...

    async def on_assistant_text_start(self, turn_id: str) -> None: ...

    async def on_assistant_text_end(
        self,
        turn_id: str,
        content: str,
        interrupted: bool,
    ) -> None: ...