"""Hermes templating engine for structured prefix generation."""

# TODO: import ansible core filters.
# TODO: Move this out of the hermes wheel and back into the tokenizer core.

import logging
import os
from logging import LoggerAdapter
from typing import NamedTuple

from chartok import heather
from constants import BOM
from data import build_template_data
from decorators import monitor
from exceptions import (
    InvalidPromptError,
    MessageStartIdxNotFound,
    TokenLimitExceededError,
)
from metrics import Metrics
from prompt_poet import AbstractTokenizer, Prompt, PromptPart
from structured_prefix import StructuredPrefix

MESSAGE_PART_PREFIX = "message_"
PINNED_MESSAGE_PART_PREFIX = "pinned_message_"
DEFAULT_TEMPLATE_NAME = "production_raw.yml.j2"
DEFAULT_TEMPLATE_DIR = "templates"

HERE = os.path.abspath(os.path.dirname(__file__))


# TODO: cleanup the metric observations.
# TODO: move top level class declarations elsewhere.
class TokenizedContext(NamedTuple):
    """Tokenized context."""

    tokens: list[int]
    idx_after_timestamp: int
    context_num_msg: int
    truncated_num_msg: int


class MonitoredPrompt(Prompt):
    """Monitored Prompt."""

    @monitor
    def __init__(self, *args, **kwargs):
        """Initialize the monitored prompt."""
        super().__init__(*args, **kwargs)

    @monitor
    def tokenize(self, *args, **kwargs):
        """Tokenize the prompt."""
        return super().tokenize(*args, **kwargs)

    @monitor
    def truncate(self, *args, **kwargs):
        """Truncate the prompt."""
        return super().truncate(*args, **kwargs)

    @monitor
    def lreplace_at(self, *args, **kwargs):
        """Replace substring of part at a specific index."""
        return super().lreplace_at(*args, **kwargs)

    @monitor
    def _render_parts(self, *args, **kwargs):
        """Render parts of the prompt."""
        return super()._render_parts(*args, **kwargs)


def _build_tokenized_context(prompt: Prompt) -> TokenizedContext:
    """Build tokenized context."""
    prompt_parts: list[PromptPart] = prompt.parts
    pretruncation_prompt_parts: list[PromptPart] = prompt.pretruncation_parts
    first_part = prompt_parts[0]

    num_prompt_messages = sum(
        1 for part in prompt_parts if part.name.startswith(MESSAGE_PART_PREFIX)
    )
    num_pretruncation_prompt_messages = sum(
        1
        for part in pretruncation_prompt_parts
        if part.name.startswith(MESSAGE_PART_PREFIX)
    )
    _observe_prompt_message_metrics(
        prompt=prompt,
        num_prompt_messages=num_prompt_messages,
        num_pretruncation_prompt_messages=num_pretruncation_prompt_messages,
    )

    return TokenizedContext(
        tokens=prompt.tokens,
        idx_after_timestamp=len(first_part.tokens) + 1,
        context_num_msg=num_prompt_messages,
        truncated_num_msg=num_pretruncation_prompt_messages - num_prompt_messages,
    )


def _observe_prompt_message_metrics(
    prompt: Prompt, num_prompt_messages: int, num_pretruncation_prompt_messages: int
):
    """Observe prompt message metrics."""
    Metrics().MESSAGE_METRICS.labels(metric_type="total_rendered_messages").observe(
        num_pretruncation_prompt_messages
    )
    Metrics().MESSAGE_METRICS.labels(metric_type="truncated_messages").observe(
        num_pretruncation_prompt_messages - num_prompt_messages
    )
    Metrics().MESSAGE_METRICS.labels(metric_type="total_active_messages").observe(
        num_prompt_messages
    )
    num_pinned_messages = sum(
        1 for part in prompt.parts if part.name.startswith(PINNED_MESSAGE_PART_PREFIX)
    )

    if num_pinned_messages > 0:
        Metrics().MESSAGE_METRICS.labels(
            metric_type="truncated_messages_with_pinned"
        ).observe(num_pretruncation_prompt_messages - num_prompt_messages)
        Metrics().MESSAGE_METRICS.labels(
            metric_type="total_active_messages_with_pinned"
        ).observe(num_prompt_messages)


def _observe_prompt_token_metrics(prompt: Prompt):
    """Observe prompt metrics."""
    Metrics().TOKEN_METRICS.labels(metric_type="total_tokens").observe(
        len(prompt.pretruncation_tokens)
    )
    Metrics().TOKEN_METRICS.labels(metric_type="truncated_tokens").observe(
        len(prompt.pretruncation_tokens) - len(prompt.tokens)
    )
    Metrics().TOKEN_METRICS.labels(metric_type="active_tokens").observe(
        len(prompt.tokens)
    )

    num_pinned_message_tokens = sum(
        len(part.tokens)
        for part in prompt.parts
        if part.name.startswith(PINNED_MESSAGE_PART_PREFIX)
    )
    if num_pinned_message_tokens > 0:
        Metrics().TOKEN_METRICS.labels(
            metric_type="total_pinned_message_tokens"
        ).observe(num_pinned_message_tokens)


def _validate_prompt(prompt: Prompt):
    """Raises if the prompt is invalid by certain heuristics."""
    # Cache prompt tokens and prompt string as it can be expensive to compute.
    prompt_tokens: list[int] = prompt.tokens
    prompt_parts: list[PromptPart] = prompt.parts
    pretruncation_prompt_parts: list[PromptPart] = prompt.pretruncation_parts
    num_pretruncation_prompt_messages = sum(
        1
        for part in pretruncation_prompt_parts
        if part.name.startswith(MESSAGE_PART_PREFIX)
    )

    if len(prompt_parts) < 2:
        raise InvalidPromptError(
            f"Prompt should have at least 2 parts: {len(prompt_parts)=}"
        )

    if prompt_parts[1].raw_string.startswith(BOM):
        raise InvalidPromptError(
            f"First message should NOT contain {BOM=}: {prompt_parts[1]=}"
        )

    if prompt.token_limit > 0 and len(prompt_tokens) > prompt.token_limit:
        raise TokenLimitExceededError(
            f"Token limit exceeded: {len(prompt_tokens)=} {prompt.token_limit=}"
        )

    if num_pretruncation_prompt_messages < 1:
        raise MessageStartIdxNotFound(
            f"Message start index not found: {num_pretruncation_prompt_messages=}"
        )


@monitor
def build_structured_prefix(
    logger: LoggerAdapter,
    structured_prefix: StructuredPrefix,
    close_last_message: bool = False,
    truncation_step: int = None,
    stream_params: dict = None,
    tokenizer: AbstractTokenizer = None,
) -> dict:
    """Build structured prefix using Prompt Poet (PP)."""
    del close_last_message  # TODO: support long streaming.
    # TODO: remove this and use package_name.
    template_name = structured_prefix.hermes_generation_template_name
    if template_name and "/" not in template_name:
        structured_prefix.hermes_generation_template_name = os.path.join(
            HERE, DEFAULT_TEMPLATE_DIR, template_name
        )

    prompt = build_prompt(
        logger=logger,
        structured_prefix=structured_prefix,
        truncation_step=truncation_step,
        stream_params=stream_params,
        tokenizer=tokenizer,
    )

    tokenized_context: TokenizedContext = _build_tokenized_context(prompt)
    first_part: PromptPart = prompt.parts[0]
    last_part: PromptPart = prompt.parts[-1]

    return {
        "character_definitions": structured_prefix.character_definitions,
        "chat_history": (
            []
            if not tokenized_context.context_num_msg
            else structured_prefix.chat_history[-tokenized_context.context_num_msg :]
        ),
        "chat_hist_global_index": 0,
        "reply_prompt": last_part.raw_string,
        "space_added": True,
        "token_limit": structured_prefix.token_limit,
        "tokenized_context": tokenized_context,
        "timestamp": first_part.raw_string,
        "prompt": prompt,
    }


@monitor
def build_prompt(
    logger: LoggerAdapter,
    structured_prefix: StructuredPrefix,
    truncation_step: int = None,
    stream_params: dict = None,
    tokenizer: AbstractTokenizer = None,
) -> Prompt:
    """Build prompt using Prompt Poet (PP)."""
    if tokenizer is None:
        tokenizer = heather()
    if logger is None:
        logger = logging.getLogger(__name__)

    template_data = build_template_data(
        structured_prefix=structured_prefix,
        logger=logger,
        stream_params=stream_params,
    )

    template_path = structured_prefix.hermes_generation_template_name
    if not template_path:
        template_path = os.path.join(HERE, DEFAULT_TEMPLATE_DIR, DEFAULT_TEMPLATE_NAME)

    prompt = MonitoredPrompt(
        template_data=template_data,
        logger=logger,
        template_path=template_path,
        raw_template=structured_prefix.hermes_generation_template,
        tokenizer=tokenizer,
        truncation_step=truncation_step,
        token_limit=structured_prefix.token_limit,
        from_cache=True,
    )
    prompt.tokenize()
    prompt.truncate()
    prompt.lreplace_at(old=BOM, new=" ", index=1)
    _validate_prompt(prompt)
    if logger:
        _observe_prompt_token_metrics(prompt)
        Metrics().HERMES_TEMPLATE_NAME.labels(
            template_name=prompt.template_name,
            template_dir=prompt.template_dir,
            source="hermes",
        ).inc()

    return prompt
