"""Assistant: a request answered by the instance's language model.

The request arrives on ``prompt``, typically what a wren node could not do.
Every output but ``reply`` and ``actions`` is a target the model may set,
offered as one tool named after its port and typed by its message, so the
wiring is the allowlist, as for wren: the MCP authoring tools are never
offered. Every other input is read as context ("outside_temperature: 7.5").

An answer goes out on ``reply`` a sentence at a time while the model writes
it, so a ``tts`` node speaks the first before the last is written. A change is
confirmed from the values set ("Kitchen light off."), without asking the model
a second time.
"""

from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import json
import logging
import re
import threading
from collections.abc import Coroutine, Iterable
from concurrent.futures import Future
from typing import TYPE_CHECKING, Any

import httpx
from pydantic import BaseModel, ConfigDict, Field

from fluksio.ai import agent, client
from fluksio.ai import config as ai_config
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes.base import Node

if TYPE_CHECKING:
    from fastapi import FastAPI

logger = logging.getLogger(__name__)

REPLY = "reply"
ACTIONS = "actions"
#: An optional input: true warms the provider while the request is spoken.
LISTENING = "listening"

#: Said when there is no provider, or it does not answer.
UNAVAILABLE = "The assistant isn't available right now."
#: Said when the model ends a turn without a word or a change.
CANT = "Sorry, I can't do that."

#: A server that is not there is given up on after a few seconds, while a
#: model that has to be loaded first takes about ten.
TIMEOUT = httpx.Timeout(15.0, connect=3.0)

#: A target's type, as its tool's parameter.
_SCHEMAS: dict[DType, dict[str, str]] = {
    DType.BOOL: {"type": "boolean"},
    DType.INT: {"type": "integer"},
    DType.FLOAT: {"type": "number"},
    DType.STR: {"type": "string"},
}

_SYSTEM = """You are the voice of a home: someone spoke to a panel on the wall, \
and you answer out loud.

- To change something, call its tool without a word: the change is confirmed \
for you. The tools are all you can change.
- If no tool can do what was asked, or you did not understand it, say so in a \
few words, such as "Sorry, I can't do that."
- Answer questions from the readings below.
- Speak plainly: one or two short sentences, no lists, no markdown, no emoji.
"""

_END = re.compile(r"(?<=[.!?])\s+")


def phrase(port: str, value: Any, word: str | None = None) -> str:
    """A value set, as the reply says it: "kitchen light off", "canopy open"."""
    label = port.replace("_", " ")
    if word:
        return f"{label} {word}"
    if isinstance(value, bool):
        return f"{label} {'on' if value else 'off'}"
    return (
        f"{label} to {value:g}" if isinstance(value, float) else f"{label} to {value}"
    )


def cut(text: str) -> tuple[list[str], str]:
    """The sentences ``text`` has finished, and the one still being written."""
    *done, rest = _END.split(text)
    return [s.strip() for s in done if s.strip()], rest


class AssistantNode(Node):
    """Answer the request on ``prompt``, acting only on the targets wired to it.

    Inputs: ``prompt`` (the request, which wakes the node), ``listening``
    (optional; turning true opens the connection and loads the model while
    the request is still spoken) and any other, read as context.

    Outputs: ``reply``, the answer a sentence at a time, ``actions``, a list
    of records, one per value set, and any number of targets, each a tool.
    When the provider is not set up or does not answer, the reply says so.
    """

    class Params(BaseModel):
        model_config = ConfigDict(extra="allow")

        model: str = Field(
            default="", description="The model to ask. Empty asks the instance's."
        )
        extra_body: dict[str, Any] = Field(
            default_factory=lambda: {
                "temperature": 0,
                "reasoning_effort": "none",
                "keep_alive": "30m",
            },
            description=(
                "Added to every request as it is. The default answers a request "
                "the same way each time, turns thinking off and keeps the model "
                "loaded for 30 minutes on Ollama; drop what a provider refuses."
            ),
        )
        values: dict[str, dict[str, str]] = Field(
            default_factory=dict,
            description=(
                "For a str target, by port: the only values it takes, each with "
                'what the reply says for it, as {"UP": "open", "DOWN": "closed"}.'
            ),
        )

    __slots__ = ("cfg", "_listening", "_lock", "_loop")

    # A tool call is a command to a device, as wren's are.
    idempotent = False

    def __init__(
        self,
        requires: MessageSpec | Iterable[MessageSpec] = (),
        provides: MessageSpec | Iterable[MessageSpec] = (),
        params: dict[str, Any] | None = None,
        name: str | None = None,
    ):
        self.cfg = self.Params.model_validate(params or {})
        self._lock = threading.Lock()
        #: What ``listening`` held at the last run, to tell its change apart.
        self._listening: bool | None = None
        self._loop: asyncio.AbstractEventLoop | None = None
        inputs = self._normalize_ports(requires)
        outputs = self._normalize_ports(provides)
        typed = {"prompt": DType.STR, LISTENING: DType.BOOL}
        for spec in inputs:
            if spec.port in typed and spec.dtype is not typed[spec.port]:
                raise ValueError(f"'{spec.port}' is a {typed[spec.port].value}.")
        if not any(spec.port == "prompt" for spec in inputs):
            raise ValueError("assistant needs a 'prompt' input: the request.")
        for spec in outputs:
            if spec.port == REPLY and spec.dtype is not DType.STR:
                raise ValueError(f"'{REPLY}' is a str.")
            if spec.port == ACTIONS and (
                spec.dtype is not DType.LIST or spec.item_dtype is not DType.RECORD
            ):
                raise ValueError(f"'{ACTIONS}' is a list of records.")
            if spec.port not in (REPLY, ACTIONS) and spec.dtype not in _SCHEMAS:
                raise ValueError(
                    f"'{spec.port}' is {spec.dtype.value}; a target is bool, int, "
                    "float or str."
                )
        super().__init__(
            self._answer,
            requires=[
                spec.model_copy(update={"trigger": spec.port in ("prompt", LISTENING)})
                for spec in inputs
            ],
            provides=outputs,
            params=params,
            name=name or "assistant",
        )

    async def start(self, app: FastAPI | None = None) -> None:
        self._loop = asyncio.get_running_loop()
        names = [spec.name for spec in self.input_ports if spec.port == LISTENING]
        if names and self._pipeline is not None:
            held = await asyncio.to_thread(self._pipeline.state.get_present, names)
            self._listening = held.get(names[0])

    def _targets(self) -> list[MessageSpec]:
        return [
            spec
            for spec in self.output_ports
            if spec.name and spec.port not in (REPLY, ACTIONS)
        ]

    def _tool(self, spec: MessageSpec) -> dict[str, Any]:
        label = spec.port.replace("_", " ")
        schema: dict[str, Any] = _SCHEMAS[spec.dtype]
        description = f"Set the {label}."
        words = self.cfg.values.get(spec.port)
        if words:
            # A code such as UP means nothing to the model until it is told.
            schema = {**schema, "enum": list(words)}
            meant = ", ".join(f'"{v}" means {label} {w}' for v, w in words.items())
            description = f"Set the {label}: {meant}."
        return {
            "type": "function",
            "function": {
                "name": spec.port,
                "description": description,
                "parameters": {
                    "type": "object",
                    "properties": {"value": schema},
                    "required": ["value"],
                },
            },
        }

    def _tools(self) -> list[dict[str, Any]]:
        return [self._tool(spec) for spec in self._targets()]

    def _system(self, context: dict[str, Any]) -> str:
        """The instructions, then what the inputs read and the targets hold."""
        targets = self._targets()
        held = (
            self._pipeline.state.get_present([spec.name for spec in targets])
            if self._pipeline is not None
            else {}
        )
        lines = [_SYSTEM, "Readings:"]
        lines += [f"{port}: {json.dumps(value)}" for port, value in context.items()]
        lines += ["", "What the tools set, now:"]
        lines += [
            f"{spec.port}: {json.dumps(held[spec.name])}"
            for spec in targets
            if spec.name in held
        ]
        return "\n".join(lines)

    def _config(self) -> ai_config.AIConfig | None:
        config = ai_config.load()
        if config is None or not config.ready:
            return None
        return dataclasses.replace(config, model=self.cfg.model or config.model)

    def _call(self, work: Coroutine[Any, Any, Any]) -> Future[Any]:
        """Run on the engine's loop, where the provider's connection is kept."""
        if self._loop is None:
            work.close()
            raise RuntimeError(f"{self.name} only runs in a live flow")
        return asyncio.run_coroutine_threadsafe(work, self._loop)

    def _answer(
        self,
        params: dict[str, Any],
        prompt: str = "",
        listening: bool | None = None,
        **context: Any,
    ) -> dict[str, Any] | None:
        with self._lock:
            was, self._listening = self._listening, listening
        config = self._config()
        if listening != was:
            # Not a request: the device started or stopped listening.
            if listening and config is not None:
                self._call(self._warm(config))
            return None
        if not prompt.strip():
            return None
        actions: list[dict[str, Any]] = []
        problem = self._call(
            self._turn(config, prompt, self._system(context), actions)
        ).result()
        if problem:
            self.report_health("degraded", problem)
        elif self._health != "ok":
            self.report_health("ok")
        wired = {spec.port for spec in self.output_ports if spec.name}
        return {ACTIONS: actions} if ACTIONS in wired else None

    async def _warm(self, config: ai_config.AIConfig) -> None:
        """A request for one token: the connection opens and the model loads.

        Nothing more is worth sending ahead: a model with a sliding window,
        such as Gemma, reads the whole prompt again on every request anyway.
        """
        with contextlib.suppress(Exception):
            async for _ in client.stream_completion(
                api_base=config.api_base,
                api_key=config.api_key,
                model=config.model,
                messages=[{"role": "user", "content": "Hello."}],
                tools=[],
                timeout=60.0,
                extra={**self.cfg.extra_body, "max_tokens": 1},
            ):
                pass

    async def _turn(
        self,
        config: ai_config.AIConfig | None,
        prompt: str,
        system: str,
        actions: list[dict[str, Any]],
    ) -> str:
        """Ask, act and say. Returns what went wrong, or an empty string."""
        said = False

        async def say(text: str) -> None:
            nonlocal said
            said = True
            await asyncio.to_thread(self._say, text)

        if config is None:
            await say(UNAVAILABLE)
            return "no AI provider is set up"
        rest, problem = "", ""
        try:
            async for event in agent.run_turn(
                config,
                [],
                prompt,
                tools=self._tools(),
                system=system,
                call_tool=self._caller(actions),
                extra=self.cfg.extra_body,
                timeout=TIMEOUT,
                answer_tools=False,
            ):
                if event["type"] == "delta":
                    done, rest = cut(rest + event["content"])
                    for sentence in done:
                        await say(sentence)
                elif event["type"] == "message" and rest.strip():
                    await say(rest.strip())
                    rest = ""
                elif event["type"] == "error":
                    problem = str(event["detail"])
        except Exception as exc:  # noqa: BLE001 - said, then reported as the health
            problem = f"{type(exc).__name__}: {exc}"
        if problem:
            logger.warning("'%s' got no answer: %s", self.id, problem)
        if not said:
            await say(UNAVAILABLE if problem else self._confirm(actions) or CANT)
        return problem

    def _confirm(self, actions: list[dict[str, Any]]) -> str:
        """The values set, as one sentence: "Kitchen light off, setpoint to 22."."""
        ports = {spec.name: spec.port for spec in self._targets()}
        said = []
        for action in actions:
            port, value = ports[action["message"]], action["value"]
            said.append(phrase(port, value, self.cfg.values.get(port, {}).get(value)))
        text = ", ".join(said)
        return text[:1].upper() + text[1:] + "." if text else ""

    def _caller(self, actions: list[dict[str, Any]]) -> agent.ToolCaller:
        """Set the target a call names, and record it in ``actions``."""
        specs = {spec.port: spec for spec in self._targets()}

        async def call(name: str, arguments: dict[str, Any]) -> str:
            spec = specs.get(name)
            if spec is None:
                return json.dumps({"error": f"There is no tool named {name}."})
            value = arguments.get("value")
            if spec.dtype is DType.FLOAT and type(value) is int:
                value = float(value)
            try:
                spec.check(value)
            except TypeError as exc:
                return json.dumps({"error": str(exc)})
            words = self.cfg.values.get(name)
            if words and value not in words:
                # A device takes its own codes only: a motor would run "OPEN"
                # as a direction it has never heard of.
                return json.dumps({"error": f"{name} is one of {', '.join(words)}."})
            await asyncio.to_thread(self.emit, {name: value})
            actions.append(
                {
                    "message": spec.name,
                    "value": value,
                    "applied": True,
                    "delay_s": None,
                    "call": f"{name}({json.dumps(value)})",
                }
            )
            return json.dumps({name: value})

        return call

    def _say(self, text: str) -> None:
        if any(spec.port == REPLY and spec.name for spec in self.output_ports):
            self.emit({REPLY: text})