import gc
import os
import unittest
from collections import defaultdict
from typing import Dict, List

import ray
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

from trinity.common.config import (
    Config,
    ExperienceBufferConfig,
    FormatConfig,
    LoRAConfig,
    TasksetConfig,
    load_config,
)
from trinity.common.constants import (
    CHECKPOINT_ROOT_DIR_ENV_VAR,
    MODEL_PATH_ENV_VAR,
    PromptType,
    StorageType,
)

API_MODEL_PATH_ENV_VAR = "TRINITY_API_MODEL_PATH"
MOE_MODEL_PATH_ENV_VAR = "TRINITY_MOE_MODEL_PATH"
VLM_MODEL_PATH_ENV_VAR = "TRINITY_VLM_MODEL_PATH"
ALTERNATIVE_VLM_MODEL_PATH_ENV_VAR = "TRINITY_ALTERNATIVE_VLM_MODEL_PATH"
SFT_DATASET_PATH_ENV_VAR = "TRINITY_SFT_DATASET_PATH"


# Qwen2.5 chat template with {% generation %} mark
CHAT_TEMPLATE_QWEN2_5 = r"""
{%- if tools %}
    {{- '<|im_start|>system\n' }}
    {%- if messages[0]['role'] == 'system' %}
        {{- messages[0]['content'] }}
    {%- else %}
        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
    {%- endif %}
    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
    {%- for tool in tools %}
        {{- "\n" }}
        {{- tool | tojson }}
    {%- endfor %}
    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
{%- else %}
    {%- if messages[0]['role'] == 'system' %}
        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
    {%- else %}
        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
    {%- endif %}
{%- endif %}
{%- for message in messages %}
    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
    {%- elif (message.role == "assistant" and not message.tool_calls) %}
        {{- '<|im_start|>' + message.role + '\n'}}{% generation %}{{- message.content + '<|im_end|>' + '\n' }}{% endgeneration %}
    {%- elif message.role == "assistant" %}
        {{- '<|im_start|>' + message.role + '\n'}}{% generation %}
        {%- if message.content %}
            {{- message.content }}
        {%- endif %}
        {%- for tool_call in message.tool_calls %}
            {%- if tool_call.function is defined %}
                {%- set tool_call = tool_call.function %}
            {%- endif %}
            {{- '<tool_call>\n{"name": "' }}
            {{- tool_call.name }}
            {{- '", "arguments": ' }}
            {{- tool_call.arguments | tojson }}
            {{- '}\n</tool_call>\n' }}
        {%- endfor %}
        {{- '<|im_end|>\n' }}{% endgeneration %}
    {%- elif message.role == "tool" %}
        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
            {{- '<|im_start|>user' }}
        {%- endif %}
        {{- '\n<tool_response>\n' }}
        {{- message.content }}
        {{- '\n</tool_response>' }}
        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
            {{- '<|im_end|>\n' }}
        {%- endif %}
    {%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
    {{- '<|im_start|>assistant\n' }}
{%- endif %}
"""

# Qwen3.5 chat template with {% generation %} mark
CHAT_TEMPLATE = r"""
{%- if tools %}
    {{- '<|im_start|>system\n' }}
    {%- if messages[0].role == 'system' %}
        {{- messages[0].content + '\n\n' }}
    {%- endif %}
    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
    {%- for tool in tools %}
        {{- "\n" }}
        {{- tool | tojson }}
    {%- endfor %}
    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
{%- else %}
    {%- if messages[0].role == 'system' %}
        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
    {%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
    {%- set index = (messages|length - 1) - loop.index0 %}
    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
        {%- set ns.multi_step_tool = false %}
        {%- set ns.last_query_index = index %}
    {%- endif %}
{%- endfor %}
{%- for message in messages %}
    {%- if message.content is string %}
        {%- set content = message.content %}
    {%- else %}
        {%- set content = '' %}
    {%- endif %}
    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
    {%- elif message.role == "assistant" %}
        {%- set reasoning_content = '' %}
        {%- if message.reasoning_content is string %}
            {%- set reasoning_content = message.reasoning_content %}
        {%- else %}
            {%- if '</think>' in content %}
                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
            {%- endif %}
        {%- endif %}
        {{- '<|im_start|>' + message.role + '\n' }}{% generation %}
        {%- if loop.index0 > ns.last_query_index %}
            {%- if loop.last or (not loop.last and reasoning_content) %}
                {{- '<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
            {%- else %}
                {{- content }}
            {%- endif %}
        {%- else %}
            {{- content }}
        {%- endif %}
        {%- if message.tool_calls %}
            {%- for tool_call in message.tool_calls %}
                {%- if (loop.first and content) or (not loop.first) %}
                    {{- '\n' }}
                {%- endif %}
                {%- if tool_call.function %}
                    {%- set tool_call = tool_call.function %}
                {%- endif %}
                {{- '<tool_call>\n{"name": "' }}
                {{- tool_call.name }}
                {{- '", "arguments": ' }}
                {%- if tool_call.arguments is string %}
                    {{- tool_call.arguments }}
                {%- else %}
                    {{- tool_call.arguments | tojson }}
                {%- endif %}
                {{- '}\n</tool_call>' }}
            {%- endfor %}
        {%- endif %}
        {{- '<|im_end|>\n' }}{% endgeneration %}
    {%- elif message.role == "tool" %}
        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
            {{- '<|im_start|>user' }}
        {%- endif %}
        {{- '\n<tool_response>\n' }}
        {{- content }}
        {{- '\n</tool_response>' }}
        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
            {{- '<|im_end|>\n' }}
        {%- endif %}
    {%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
    {{- '<|im_start|>assistant\n' }}
    {%- if enable_thinking is defined and enable_thinking is false %}
        {{- '<think>\n\n</think>\n\n' }}
    {%- endif %}
{%- endif %}
"""


def get_template_config() -> Config:
    config_path = os.path.join(os.path.dirname(__file__), "template", "config.yaml")
    config = load_config(config_path)
    if ray.is_initialized():
        config.ray_namespace = ray.get_runtime_context().namespace
    else:
        config.ray_namespace = "trinity_unittest"
    return config


def get_model_path() -> str:
    path = os.environ.get(MODEL_PATH_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {MODEL_PATH_ENV_VAR}=<your_model_dir>` before running this test."
        )
    return path


def get_api_model_path() -> str:
    path = os.environ.get(API_MODEL_PATH_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {API_MODEL_PATH_ENV_VAR}=<your_api_model_checkpoint_dir>` before running this test."
        )
    return path


def get_moe_model_path() -> str:
    path = os.environ.get(MOE_MODEL_PATH_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {MOE_MODEL_PATH_ENV_VAR}=<your_moe_model_checkpoint_dir>` before running this test."
        )
    return path


def get_checkpoint_path() -> str:
    path = os.environ.get(CHECKPOINT_ROOT_DIR_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {CHECKPOINT_ROOT_DIR_ENV_VAR}=<your_checkpoint_dir>` before running this test."
        )
    return path


def get_vision_language_model_path() -> str:
    path = os.environ.get(VLM_MODEL_PATH_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {VLM_MODEL_PATH_ENV_VAR}=<your_model_dir>` before running this test."
        )
    return path


def get_alternative_vision_language_model_path() -> str:
    path = os.environ.get(ALTERNATIVE_VLM_MODEL_PATH_ENV_VAR)
    if not path:
        raise EnvironmentError(
            f"Please set `export {ALTERNATIVE_VLM_MODEL_PATH_ENV_VAR}=<your_model_dir>` before running this test."
        )
    return path


def get_lora_config() -> LoRAConfig:
    return LoRAConfig(name="lora", lora_rank=16, lora_alpha=16)


def get_unittest_dataset_config(dataset_name: str = "countdown", split: str = "train"):
    if dataset_name == "countdown" or dataset_name == "copy_countdown":
        # Countdown dataset with 17 samples
        return TasksetConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "countdown"),
            split=split,
            enable_progress_bar=False,
            format=FormatConfig(
                prompt_key="question",
                response_key="answer",
            ),
            default_workflow_type="math_workflow",
            default_reward_fn_type="countdown_reward",
        )
    elif dataset_name in {"eval_short", "eval_long"}:
        # Eval_short dataset with 2 samples, eval_long dataset with 8 samples
        return TasksetConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", dataset_name),
            split="test",
            format=FormatConfig(
                prompt_key="question",
                response_key="answer",
            ),
            default_workflow_type="math_workflow",
            default_reward_fn_type="math_reward",
        )
    elif dataset_name == "gsm8k":
        # GSM8K dataset with 16 samples
        return TasksetConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "gsm8k"),
            split="train",
            format=FormatConfig(
                prompt_key="question",
                response_key="answer",
            ),
            default_workflow_type="math_workflow",
            default_reward_fn_type="math_reward",
        )
    elif dataset_name == "gsm8k_ruler":
        return TasksetConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "gsm8k"),
            split="train",
            format=FormatConfig(
                prompt_key="question",
                response_key="answer",
            ),
            default_workflow_type="math_ruler_workflow",
        )
    elif dataset_name == "sft_for_gsm8k":
        # SFT dataset with 8 samples
        return ExperienceBufferConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "sft_for_gsm8k"),
            split="train",
            storage_type=StorageType.FILE.value,
            schema_type="sft",
            format=FormatConfig(
                prompt_type=PromptType.PLAINTEXT,
                prompt_key="prompt",
                response_key="response",
            ),
        )
    elif dataset_name == "sft_with_tools":
        # SFT_with_tools dataset with 4 samples
        return ExperienceBufferConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "sft_with_tools"),
            split="train",
            storage_type=StorageType.FILE.value,
            format=FormatConfig(
                prompt_type=PromptType.MESSAGES,
                messages_key="messages",
                tools_key="tools",
                enable_concatenated_multi_turn=True,
            ),
        )
    elif dataset_name == "dpo":
        # HumanLike DPO dataset with 17 samples
        return ExperienceBufferConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "human_like"),
            split="train",
            storage_type=StorageType.FILE.value,
            schema_type="dpo",
            format=FormatConfig(
                prompt_type=PromptType.PLAINTEXT,
                prompt_key="prompt",
                chosen_key="chosen",
                rejected_key="rejected",
            ),
        )
    elif dataset_name == "geometry":
        # Multi-modal geometry dataset with 8 samples
        return TasksetConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "geometry"),
            split="train",
            format=FormatConfig(
                prompt_type=PromptType.PLAINTEXT,
                prompt_key="problem",
                response_key="answer",
                image_key="images",
            ),
            default_workflow_type="simple_mm_workflow",
            default_reward_fn_type="math_boxed_reward",
        )
    elif dataset_name == "geometry_sft":
        # Multi-modal geometry dataset for sft with 8 samples
        return ExperienceBufferConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "geometry"),
            split="train",
            storage_type=StorageType.FILE.value,
            format=FormatConfig(
                prompt_type=PromptType.PLAINTEXT,
                prompt_key="problem",
                response_key="answer",
                image_key="images",
            ),
        )
    elif dataset_name == "qwenpaw_sft":
        return ExperienceBufferConfig(
            name=dataset_name,
            path=os.path.join(os.path.dirname(__file__), "template", "data", "qwenpaw_sft"),
            split="train",
            storage_type=StorageType.FILE.value,
            format=FormatConfig(
                prompt_type=PromptType.MESSAGES,
                messages_key="messages",
                tools_key="tools",
                enable_concatenated_multi_turn=True,
            ),
        )
    else:
        raise ValueError(f"Unknown dataset name: {dataset_name}")


class TensorBoardParser:
    def __init__(self, log_dir: str):
        self.log_dir = log_dir
        self._event_files = self._find_event_files(log_dir)
        self._metrics = self._load_metrics()

    def _find_event_files(self, log_dir: str) -> List[str]:
        event_files = []
        for root, _, files in os.walk(log_dir):
            for f in files:
                if f.startswith("events.out.tfevents."):
                    event_files.append(os.path.join(root, f))
        return event_files

    def _load_metrics(self) -> Dict[str, Dict[int, float]]:
        metrics = defaultdict(dict)

        for event_file in self._event_files:
            ea = EventAccumulator(event_file)
            ea.Reload()
            tags = ea.Tags()["scalars"]
            for tag in tags:
                scalars = ea.Scalars(tag)
                for scalar in scalars:
                    step = scalar.step
                    value = scalar.value
                    if step not in metrics[tag] or value > metrics[tag][step]:
                        metrics[tag][step] = value
        return dict(metrics)

    def metric_exist(self, metric_name: str) -> bool:
        return metric_name in self._metrics

    def metric_min_step(self, metric_name: str) -> int:
        return min(self.metric_steps(metric_name))

    def metric_max_step(self, metric_name: str) -> int:
        return max(self.metric_steps(metric_name))

    def metric_steps(self, metric_name: str) -> List[int]:
        if not self.metric_exist(metric_name):
            raise ValueError(f"Metric '{metric_name}' does not exist.")
        return list(self._metrics[metric_name].keys())

    def metric_values(self, metric_name: str) -> List:
        if not self.metric_exist(metric_name):
            raise ValueError(f"Metric '{metric_name}' does not exist.")
        return list(self._metrics[metric_name].values())

    def metric_list(self, metric_prefix: str) -> List[str]:
        return [name for name in self._metrics if name.startswith(metric_prefix)]


class RayCleanupPlugin:
    @classmethod
    def _cleanup_ray_data_state(cls):
        """clean up the global states of Ray Data"""
        try:
            # reset execution context
            if hasattr(ray.data._internal.execution.streaming_executor, "_execution_context"):
                ray.data._internal.execution.streaming_executor._execution_context = None

            # trigger gc.collect() on all workers in the cluster
            ray._private.internal_api.global_gc()

            # clean up stats manager
            from ray.data._internal.stats import StatsManager

            if hasattr(StatsManager, "_instance"):
                StatsManager._instance = None

        except Exception:
            pass


class RayUnittestBase(unittest.TestCase, RayCleanupPlugin):
    @classmethod
    def setUpClass(cls):
        ray.init(ignore_reinit_error=True, namespace="trinity_unittest")

        # erase existing resources
        cls._cleanup_ray_data_state()
        gc.collect()

    @classmethod
    def tearDownClass(cls):
        ray.shutdown(_exiting_interpreter=True)


class RayUnittestBaseAsync(unittest.IsolatedAsyncioTestCase, RayCleanupPlugin):
    @classmethod
    def setUpClass(cls):
        ray.init(ignore_reinit_error=True, namespace="trinity_unittest")

        # erase existing resources
        cls._cleanup_ray_data_state()
        gc.collect()

    @classmethod
    def tearDownClass(cls):
        ray.shutdown(_exiting_interpreter=True)
