"""
Use viz_sft_dataset to visualize the output of different renderers. E.g.,
    python -m ttt_discover.supervised.viz_sft_dataset dataset_path=Tulu3Builder renderer_name=role_colon
"""

import json
import logging
import re
from datetime import datetime
from enum import StrEnum
from typing import NotRequired, TypedDict, Literal, Protocol, TYPE_CHECKING, Any, TypeAlias
from functools import cache

import tinker
import torch
import pydantic

from functools import cache
from typing import TYPE_CHECKING, Any, TypeAlias

if TYPE_CHECKING:
    # this import takes a few seconds, so avoid it on the module import when possible
    from transformers.tokenization_utils import PreTrainedTokenizer

    Tokenizer: TypeAlias = PreTrainedTokenizer
else:
    # make it importable from other files as a type in runtime
    Tokenizer: TypeAlias = Any

logger = logging.getLogger(__name__)

# Tool types are based on kosong (https://github.com/MoonshotAI/kosong).


@cache
def get_tokenizer(model_name: str) -> Tokenizer:
    import os
    from transformers.models.auto.tokenization_auto import AutoTokenizer

    # If it's a local path, use it directly
    if os.path.isdir(model_name):
        return AutoTokenizer.from_pretrained(model_name, use_fast=True, local_files_only=True)

    # Avoid gating of Llama 3 models:
    if model_name.startswith("meta-llama/Llama-3"):
        model_name = "thinkingmachineslabinc/meta-llama-3-tokenizer"

    kwargs: dict[str, Any] = {}
    if model_name == "moonshotai/Kimi-K2-Thinking":
        kwargs["trust_remote_code"] = True
        kwargs["revision"] = "612681931a8c906ddb349f8ad0f582cb552189cd"

    return AutoTokenizer.from_pretrained(model_name, use_fast=True, **kwargs)


class StrictBase(pydantic.BaseModel):
    """
    Pydantic base class that's immutable and doesn't silently ignore extra fields.
    """

    model_config = pydantic.ConfigDict(frozen=True, extra="forbid")

    def __str__(self) -> str:
        return repr(self)


class TextPart(TypedDict):
    """
    Container for a text part in a multimodal message.

    Args:

    type: Literal['text']
        The type of the content part, which must be text in this case.
    text: str
        The string content of the content part.
    """

    type: Literal["text"]
    text: str


# Container for a part of a multimodal message content
ContentPart = TextPart


# NOTE: we use a broad type definition for the role to be flexible
# Common roles are "user", "assistant", "system", "tool"
Role = str

# Content is a string or a list of parts
Content = str | list[ContentPart]


class Message(TypedDict):
    """
    Container for a single turn in a multi-turn conversation.

    Args:

    role: Role
        String that denotes the source of the message, typically system, user, assistant, and tool.
    content: Content
        Content of the message, can be a string, or a list of ContentPart.
    tool_calls: NotRequired[list[ToolCall]]
        Optional sequence of tool calls generated by the model.
    thinking: NotRequired[str]
        Optional thinking produced by the model before its final response.
    trainable: NotRequired[bool]
        Optional indicator whether this message should contribute to the training loss.

    """

    role: Role
    content: Content

    thinking: NotRequired[str]
    trainable: NotRequired[bool]
    name: NotRequired[str]


def ensure_text(content: Content) -> str:
    """
    Assert that content is text-only and return it as a string.

    Raises ValueError if content contains images or multiple parts.
    Use this to validate that message content is text-only before
    processing it in code paths that don't support multimodal content.
    """
    if isinstance(content, str):
        return content
    if len(content) == 1 and content[0]["type"] == "text":
        return content[0]["text"]
    raise ValueError(f"Expected text content, got multimodal content with {len(content)} parts")


class RenderedMessage(TypedDict):
    """
    Container for parts of a rendered message, for masking.

    Args:

    prefix: NotRequired[tinker.EncodedTextChunk]
        Message header that typically includes the speaker's role in the conversation.
    content: list[tinker.ModelInputChunk]
        Inner parts of the message that may include spans of image and text.
    suffix: NotRequired[tinker.EncodedTextChunk]
        Message header that typically includes the turn stop token.

    """

    prefix: NotRequired[tinker.EncodedTextChunk]
    content: list[tinker.ModelInputChunk]
    suffix: NotRequired[tinker.EncodedTextChunk]


class TrainOnWhat(StrEnum):
    LAST_ASSISTANT_MESSAGE = "last_assistant_message"
    ALL_ASSISTANT_MESSAGES = "all_assistant_messages"
    ALL_MESSAGES = "all_messages"
    ALL_TOKENS = "all_tokens"
    ALL_USER_AND_SYSTEM_MESSAGES = "all_user_and_system_messages"
    CUSTOMIZED = "customized"


class Renderer(Protocol):
    """
    Render a message list into training and sampling prompts for language models.
    """

    tokenizer: Tokenizer

    def __init__(self, tokenizer: Tokenizer):
        self.tokenizer = tokenizer

    def _preprocess_message_parts(self, message: Message) -> list[TextPart]:
        return (
            message["content"]
            if isinstance(message["content"], list)
            else [TextPart(type="text", text=message["content"])]
        )

    @property
    def _bos_tokens(self) -> list[int]:
        return []

    def get_stop_sequences(self) -> list[str] | list[int]:
        raise NotImplementedError

    def render_message(self, idx: int, message: Message, is_last: bool = False) -> RenderedMessage:
        raise NotImplementedError

    def parse_response(self, response: list[int]) -> tuple[Message, bool]:
        raise NotImplementedError

    def build_generation_prompt(
        self, messages: list[Message], role: Role = "assistant", prefill: str | None = None
    ) -> tinker.ModelInput:
        """
        Generates tokens for sampling from the model.

        Args:
            messages: a list of messages to render.
            role: the role of the partial message to be completed.
            prefill: an optional string to prefill in the model's generation.
        """

        chunks: list[tinker.types.ModelInputChunk] = []
        if self._bos_tokens:
            chunks.append(tinker.types.EncodedTextChunk(tokens=self._bos_tokens))
        for idx, message in enumerate(messages):
            rendered_message = self.render_message(idx, message)
            ob_chunk = rendered_message.get("prefix")
            action_chunks = rendered_message["content"]
            if ob_chunk:
                chunks.append(ob_chunk)
            chunks.extend([x for x in action_chunks if x])
        new_partial_message = Message(role=role, content="")
        rendered_message = self.render_message(len(messages), new_partial_message)
        ob_chunk = rendered_message.get("prefix")
        if ob_chunk:
            chunks.append(ob_chunk)
        if prefill:
            chunks.append(
                tinker.types.EncodedTextChunk(
                    tokens=self.tokenizer.encode(prefill, add_special_tokens=False)
                )
            )
        return tinker.ModelInput(chunks=chunks)

    def build_supervised_example(
        self,
        messages: list[Message],
        train_on_what: TrainOnWhat = TrainOnWhat.LAST_ASSISTANT_MESSAGE,
    ) -> tuple[tinker.ModelInput, torch.Tensor]:
        """
        Generates tokens and weights (for SFT) in the most standard way; by concatenating
        together tokens and weights for each message.

        Args:
            messages: a list of messages to render.
            train_on_what: an enum that controls how the weights are assigned to the tokens.
                - TrainOnWhat.LAST_ASSISTANT_MESSAGE: only the last assistant message is used for training
                - TrainOnWhat.ALL_ASSISTANT_MESSAGES: all assistant messages are used for training
                - TrainOnWhat.ALL_MESSAGES: all messages are used for training
                - TrainOnWhat.ALL_TOKENS: all tokens are used for training
                - TrainOnWhat.ALL_USER_AND_SYSTEM_MESSAGES: all user and system messages are used for training
                - TrainOnWhat.CUSTOMIZED: each message has a trainable field, and the weights are assigned based on the trainable field

        Returns:
            A tuple of two tensors:
                - model_input: the tinker ModelInput for your model
                - weights: a tensor of weights
        """

        model_input_chunks_weights: list[tuple[tinker.types.ModelInputChunk, float]] = []
        if self._bos_tokens:
            model_input_chunks_weights.append(
                (tinker.types.EncodedTextChunk(tokens=self._bos_tokens), 0.0)
            )

        for idx, message in enumerate(messages):
            if train_on_what == TrainOnWhat.CUSTOMIZED:
                assert "trainable" in message, (
                    "When using CUSTOMIZED train_on_what, each message must have a trainable field: True if loss is applied on this message, False otherwise"
                )
            else:
                assert "trainable" not in message, (
                    "When using non-CUSTOMIZED train_on_what, each message must not have a trainable field. Either change train_on_what to CUSTOMIZED or remove the trainable field from the message"
                )

            is_last_message = idx == len(messages) - 1
            is_assistant = message["role"] == "assistant"
            is_user_or_system = message["role"] in ["user", "system"]

            # only apply weight to observation part if train_on_what is ALL_TOKENS
            rendered_message = self.render_message(idx, message, is_last=is_last_message)
            ob_part = rendered_message.get("prefix")
            action_parts = rendered_message.get("content")
            action_tail = rendered_message.get("suffix")

            ob_weight = int(train_on_what == TrainOnWhat.ALL_TOKENS)
            if ob_part:
                model_input_chunks_weights += [(ob_part, ob_weight)]

            match train_on_what:
                case TrainOnWhat.LAST_ASSISTANT_MESSAGE:
                    action_has_weight = is_last_message and is_assistant
                case TrainOnWhat.ALL_ASSISTANT_MESSAGES:
                    action_has_weight = is_assistant
                case TrainOnWhat.ALL_MESSAGES:
                    action_has_weight = True
                case TrainOnWhat.ALL_TOKENS:
                    action_has_weight = True
                case TrainOnWhat.ALL_USER_AND_SYSTEM_MESSAGES:
                    action_has_weight = is_user_or_system
                case TrainOnWhat.CUSTOMIZED:
                    action_has_weight = message.get("trainable", False)
                case _:
                    raise ValueError(f"Unknown train_on_what: {train_on_what}")

            model_input_chunks_weights += [
                (action_part, int(action_has_weight)) for action_part in action_parts if action_part
            ]

            # action tail is effectively the stop_token and the start token for the next turn
            # e.g. \n\nUser:
            if is_last_message and action_tail:
                model_input_chunks_weights += [(action_tail, int(action_has_weight))]

        weights_data = [w for chunk, w in model_input_chunks_weights for _ in range(chunk.length)]
        weights_tensor = torch.tensor(weights_data)

        model_input_chunks = [chunk for chunk, _ in model_input_chunks_weights]
        return tinker.ModelInput(chunks=model_input_chunks), weights_tensor


def parse_response_for_stop_token(
    response: list[int], tokenizer: Tokenizer, stop_token: int
) -> tuple[Message, bool]:
    """Parse response for a single stop token.

    We expect a properly rendered response to have exactly one stop token; but it may have zero if e.g. the model
    ran out of tokens when sampling, which will incur a format error. If there are > 1, there is likely a bug in the
    sampler and we should error.
    """
    emt_count = response.count(stop_token)
    if emt_count == 0:
        str_response = tokenizer.decode(response)
        logger.debug(f"Response is not a valid assistant response: {str_response}")
        return Message(role="assistant", content=str_response), False
    elif emt_count == 1:
        str_response = tokenizer.decode(response[: response.index(stop_token)])
        return Message(role="assistant", content=str_response), True
    else:
        raise ValueError(
            f"When parsing response, expected to split into 1 or 2 pieces using stop tokens, but got {emt_count}. "
            "You probably are using the wrong stop tokens when sampling"
        )


class Qwen3Renderer(Renderer):
    """
    Format like this:
        <|im_start|>system
        You are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>
        <|im_start|>user
        What can you help me with?<|im_end|>
        <|im_start|>assistant
        <think>

        </think>
        I can help you with...<|im_end|>
    """

    def __init__(self, tokenizer: Tokenizer, strip_thinking_from_history: bool = True):
        """
        Args:
            tokenizer: The tokenizer to use for encoding.
            strip_thinking_from_history: When True (default), strips <think>...</think> blocks
                from assistant messages in multi-turn history. This matches how Qwen3 models
                were trained - they only see their own thinking during the current turn, not
                from previous turns. Set to False to preserve thinking in history (useful for
                certain RL scenarios where you want the extension property for efficiency).

        See https://tinker-docs.thinkingmachines.ai/rl/sequence-extension for details on
        how this option affects multi-turn RL compute efficiency.
        """
        super().__init__(tokenizer)
        self.strip_thinking_from_history = strip_thinking_from_history

    def render_message(self, idx: int, message: Message, is_last: bool = False) -> RenderedMessage:
        assert message.get("thinking") is None, "TODO: support CoT in Qwen3 renderer"
        assert isinstance(message["content"], str), (
            "Qwen3Renderer only supports message with string content"
        )
        maybe_newline = "\n" if idx > 0 else ""
        ob_str = f"{maybe_newline}<|im_start|>{message['role']}\n"
        ac_content = message["content"]
        if (
            self.strip_thinking_from_history
            and message["role"] == "assistant"
            and "</think>" in ac_content
        ):
            # Multi-turn conversation, we remove the thinking section from the assistant message.
            # This matches how Qwen3 models were trained - they only see their own thinking
            # during the current turn, not from previous turns.
            ac_content = ac_content.split("</think>")[1].lstrip()
        elif message["role"] == "assistant" and "<think>" not in ac_content:
            # Matching the paper, we force the assistant to start with <think>. Some SFT datasets include
            # <think> in the assistant messages, we so don't need to re-add it in those cases.
            ob_str += "<think>\n"
        
        ac_content += "<|im_end|>"
        # Action part
        prefix = tinker.types.EncodedTextChunk(
            tokens=self.tokenizer.encode(ob_str, add_special_tokens=False)
        )
        content: list[tinker.ModelInputChunk] = [
            tinker.types.EncodedTextChunk(
                tokens=self.tokenizer.encode(ac_content, add_special_tokens=False)
            )
        ]
        return RenderedMessage(prefix=prefix, content=content)

    @property
    def _end_message_token(self) -> int:
        tokens = self.tokenizer.encode("<|im_end|>", add_special_tokens=False)
        assert len(tokens) == 1, f"Expected single token for <|im_end|>, got {len(tokens)}"
        return tokens[0]

    def get_stop_sequences(self) -> list[int]:
        return [self._end_message_token]

    def parse_response(self, response: list[int]) -> tuple[Message, bool]:
        assistant_message, parse_success = parse_response_for_stop_token(
            response, self.tokenizer, self._end_message_token
        )
        if not parse_success:
            return assistant_message, False

        # Follow Qwen docs and Qwen-Agent's tool calling prompt to use <tool_call>...</tool_call> tags to wrap the tool call.
        # - https://qwen.readthedocs.io/en/latest/getting_started/concepts.html#tool-calling
        # - https://github.com/QwenLM/Qwen-Agent/blob/main/qwen_agent/llm/fncall_prompts/nous_fncall_prompt.py#L279-L282
        assert isinstance(assistant_message["content"], str)
        # TODO: WARN TOOL CALL IS NOT SUPPORTED.
        return assistant_message, True


class Qwen3InstructRenderer(Qwen3Renderer):
    """
    Renderer for Qwen3 instruct 2507 models. Unlike the earlier Qwen3 models, these models do not
    use the <think> tag at all.
    """

    def render_message(self, idx: int, message: Message, is_last: bool = False) -> RenderedMessage:
        assert message.get("thinking") is None, "CoT tokens not supported in Qwen3 instruct 2507"
        assert isinstance(message["content"], str), (
            "Qwen3InstructRenderer only supports message with string content"
        )
        maybe_newline = "\n" if idx > 0 else ""
        ob_str = f"{maybe_newline}<|im_start|>{message['role']}\n"
        ac_content = message["content"]
        ac_content += "<|im_end|>"
        # Action part
        prefix = tinker.types.EncodedTextChunk(
            tokens=self.tokenizer.encode(ob_str, add_special_tokens=False)
        )
        content: list[tinker.ModelInputChunk] = [
            tinker.types.EncodedTextChunk(
                tokens=self.tokenizer.encode(ac_content, add_special_tokens=False)
            )
        ]
        return RenderedMessage(prefix=prefix, content=content)


class GptOssRenderer(Renderer):
    """
    Format like this (no newlines between messages, last message should end with <|return|> but be replaced by <|end|> when continuing the convo):
        <|start|>system<|message|>You are ChatGPT...<|end|><|start|>user<|message|>How much is 1+1?<|end|><|start|>assistant<|channel|>final<|message|>2<|end|><|start|>
    TODO: support channels in input messages and tools
    """

    system_prompt = "<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\nCurrent date: {current_date}\n\nReasoning: {reasoning_effort}\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|>"
    use_system_prompt: bool = False
    reasoning_effort: str | None = None
    current_date: str | None = (
        None  # If use_system_prompt=True, will use the current date if this is None. Set this to a fixed date for deterministic system prompt.
    )

    def __init__(
        self,
        tokenizer: Tokenizer,
        use_system_prompt: bool = False,
        reasoning_effort: str | None = None,
        current_date: str | None = None,
    ):
        super().__init__(tokenizer)
        self.use_system_prompt = use_system_prompt
        self.reasoning_effort = reasoning_effort
        self.current_date = current_date
        assert use_system_prompt == (reasoning_effort is not None), (
            "Reasoning effort must be set iff using system prompt"
        )

    def render_message(self, idx: int, message: Message, is_last: bool = False) -> RenderedMessage:
        assert message.get("tool_calls") is None, "TODO: support tools in gpt-oss renderer"
        assert isinstance(message["content"], str), (
            "GptOssRenderer only supports message with string content"
        )
        # Observation (prompt) part
        ob_str = f"<|start|>{message['role']}"
        # Action part
        ac_str = ""
        if message["role"] == "assistant":
            # TODO: support commentary channel / tools

            # Assistant channels. See https://cookbook.openai.com/articles/openai-harmony
            thinking = message.get("thinking")
            message_content = message.get("content", "")
            assert isinstance(message_content, str), "GptOssRenderer only supports string content"

            # Analysis channel (CoT)
            if thinking:
                if is_last:
                    # Analysis channel only included in the last message. See https://cookbook.openai.com/articles/gpt-oss/handle-raw-cot
                    ac_str += f"<|channel|>analysis<|message|>{thinking}<|end|><|start|>assistant"

            # Final channel (Response Content)
            ac_str += f"<|channel|>final<|message|>{message_content}"
        else:
            assert message.get("thinking") is None, (
                "Thinking is only allowed for assistant messages"
            )
            ac_str += f"<|message|>{message['content']}"

        if not is_last:
            ac_str += "<|end|>"
        else:
            # <|return|> ends the last-message in harmony (but should be replaced by <|end|> when continuing the convo)
            ac_str += "<|return|>"

        prefix = tinker.types.EncodedTextChunk(
            tokens=self.tokenizer.encode(ob_str, add_special_tokens=False)
        )
        content: list[tinker.ModelInputChunk] = [
            tinker.types.EncodedTextChunk(
                tokens=self.tokenizer.encode(ac_str, add_special_tokens=False)
            )
        ]
        return RenderedMessage(prefix=prefix, content=content)

    def _build_system_prompt(self) -> str:
        current_date = (
            self.current_date
            if self.current_date is not None
            else datetime.now().strftime("%Y-%m-%d")
        )
        return self.system_prompt.format(
            current_date=current_date, reasoning_effort=self.reasoning_effort
        )

    @property
    def _bos_tokens(self) -> list[int]:
        tokens = []
        if self.use_system_prompt:
            tokens.extend(
                self.tokenizer.encode(self._build_system_prompt(), add_special_tokens=False)
            )
        return tokens

    @property
    def _return_token(self) -> int:
        res = self.tokenizer.encode("<|return|>", add_special_tokens=False)
        assert len(res) == 1, f"Expected single token for <|return|>, got {len(res)}"
        return res[0]

    def get_stop_sequences(self) -> list[int]:
        return [self._return_token]

    def parse_response(self, response: list[int]) -> tuple[Message, bool]:
        return parse_response_for_stop_token(response, self.tokenizer, self._return_token)


def get_renderer(
    name: str, tokenizer: Tokenizer
) -> Renderer:
    if name == "qwen3":
        return Qwen3Renderer(tokenizer)
    elif name == "qwen3_instruct":
        return Qwen3InstructRenderer(tokenizer)
    elif name == "gpt_oss_no_sysprompt":
        return GptOssRenderer(tokenizer, use_system_prompt=False)
    elif name == "gpt_oss_low_reasoning":
        return GptOssRenderer(tokenizer, use_system_prompt=True, reasoning_effort="low")
    elif name == "gpt_oss_medium_reasoning":
        return GptOssRenderer(tokenizer, use_system_prompt=True, reasoning_effort="medium")
    elif name == "gpt_oss_high_reasoning":
        return GptOssRenderer(tokenizer, use_system_prompt=True, reasoning_effort="high")
    else:
        raise ValueError(f"Unknown renderer: {name}")
