"""A stand-in model for NeMo Guardrails, so every lesson runs with no API key.

Written across lessons 4, 5, 6 and 9 of the NeMo Guardrails course. Importing
it registers two providers under the name "pretend": a chat model and an
embedding model.
"""

import re
import zlib

from nemoguardrails import LLMResponse, LLMResponseChunk, register_provider
from nemoguardrails.context import llm_call_info_var
from nemoguardrails.embeddings.providers import register_embedding_provider
from nemoguardrails.embeddings.providers.base import EmbeddingModel

HANDBOOK = {
    "a17": "Order A17 shipped on 3 March by courier.",
    "refund": "Refunds take five working days.",
    "refunds": "Refunds take five working days.",
    "delivery": "Delivery is free on orders over 40 pounds.",
}

SKIP = {"the", "and", "for", "you", "your", "with", "what", "can", "how", "this",
        "that", "its", "from", "are", "was", "about", "tell", "give", "say"}


def words(text):
    """The words worth comparing, as a set."""
    return {w for w in re.findall(r"[a-z0-9]+", str(text).lower())
            if len(w) > 1 and w not in SKIP}


def task_name():
    """Which task the runtime is asking about, or None outside a task."""
    info = llm_call_info_var.get()
    return info.task if info else None


def last_turn(prompt):
    """What the user said, in whichever shape this prompt uses."""
    found = re.findall(r'User message: "(.*)"', prompt) or re.findall(r"User: (.*)", prompt)
    return found[-1] if found else ""


def examples(prompt):
    """The (utterance, intent) pairs the runtime pasted into the prompt."""
    block = prompt.split("# This is how the user talks:")[-1]
    block = block.split("# This is the current conversation")[0]
    return re.findall(r'User message: "(.*)"\nUser intent: (.*)', block)


def handbook_answer(question):
    for word in words(question):
        if word in HANDBOOK:
            return HANDBOOK[word]
    return "I do not have that in the handbook."


def best_intent(prompt):
    """Pick the intent whose example shares the most words with this message."""
    asked = words(last_turn(prompt))
    best, score = "", 0
    for utterance, intent in examples(prompt):
        shared = len(asked & words(utterance))
        if shared > score:
            best, score = intent.strip(), shared
    return best or "ask general question"


def self_check(prompt):
    """Yes blocks, No allows. The policy comes from the prompt you wrote."""
    banned = re.search(r"Blocked topics: (.*)", prompt)
    checked = re.findall(r'Message: "(.*)"', prompt)
    text = checked[-1].lower() if checked else ""
    if banned and any(t.strip() in text for t in banned.group(1).split(",")):
        return "Yes"
    return "No"


class PretendLLM:
    """Answers the runtime's tasks without a network call."""

    def __init__(self, model="pretend-1", **kwargs):
        self._model = model
        self.replies = []
        self.seen = []
        self.prompts = []

    model_name = property(lambda self: self._model)
    provider_name = property(lambda self: "pretend")
    provider_url = property(lambda self: None)

    def say(self, text, task=None):
        """Queue one exact answer, for the next call or for one named task."""
        self.replies.append((task, text))

    def queued(self, task):
        for i, (wanted, text) in enumerate(self.replies):
            if wanted is None or wanted == task:
                self.replies.pop(i)
                return text
        return None

    def answer(self, task, prompt):
        if task == "generate_user_intent":
            return "User intent: " + best_intent(prompt)
        if task == "generate_next_steps":
            return "Bot intent: answer the question"
        if task == "generate_bot_message":
            return 'Bot message: "%s"' % handbook_answer(last_turn(prompt))
        if task and task.startswith("self_check"):
            return self_check(prompt)
        return handbook_answer(last_turn(prompt))

    async def generate_async(self, prompt, *, stop=None, **kwargs):
        text = prompt if isinstance(prompt, str) else "\n".join(m.content for m in prompt)
        task = task_name()
        self.seen.append(task)
        self.prompts.append(text)
        reply = self.queued(task)
        if reply is None:
            reply = self.answer(task, text)
        return LLMResponse(content=reply, model=self._model, finish_reason="stop")

    async def stream_async(self, prompt, *, stop=None, **kwargs):
        response = await self.generate_async(prompt, stop=stop, **kwargs)
        yield LLMResponseChunk(delta_content=response.content, model=self._model)
        yield LLMResponseChunk(model=self._model, finish_reason="stop")


class PretendEmbeddings(EmbeddingModel):
    """One slot per word, so texts sharing words end up close together."""

    engine_name = "pretend"
    WIDTH = 64

    def __init__(self, embedding_model="pretend-embed", **kwargs):
        self.model_name = embedding_model
        self.embedding_size = self.WIDTH

    def encode(self, documents):
        return [self.vector(d) for d in documents]

    async def encode_async(self, documents):
        return self.encode(documents)

    def vector(self, text):
        out = [0.0] * self.WIDTH
        for word in words(text):
            out[zlib.crc32(word.encode()) % self.WIDTH] = 1.0
        size = sum(out) ** 0.5 or 1.0
        return [v / size for v in out]


register_provider("pretend", PretendLLM)
register_embedding_provider(PretendEmbeddings, "pretend")