from __future__ import annotations

import importlib.util
import math
from pathlib import Path
import sys
import types
import unittest
from unittest import mock


PLUGIN_ROOT = Path(__file__).resolve().parents[1]


def _load_source_module(name: str, path: Path):
    spec = importlib.util.spec_from_file_location(name, path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"无法加载测试模块：{path}")
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    spec.loader.exec_module(module)
    return module


def _load_prompt_module():
    root_package = types.ModuleType("comfyui_o1key")
    root_package.__path__ = [str(PLUGIN_ROOT)]
    sys.modules["comfyui_o1key"] = root_package

    for package_name, directory in (
        ("comfyui_o1key.utils", PLUGIN_ROOT / "utils"),
        ("comfyui_o1key.nodes", PLUGIN_ROOT / "nodes"),
    ):
        package = types.ModuleType(package_name)
        package.__path__ = [str(directory)]
        sys.modules[package_name] = package

    _load_source_module(
        "comfyui_o1key.utils.image_utils",
        PLUGIN_ROOT / "utils" / "image_utils.py",
    )
    return _load_source_module(
        "comfyui_o1key.nodes.prompt_multi_function",
        PLUGIN_ROOT / "nodes" / "prompt_multi_function.py",
    )


prompt_module = _load_prompt_module()
O1keyPromptMultiFunction = prompt_module.O1keyPromptMultiFunction


def _prompts(count: int) -> str:
    return "\n---\n".join(f"提示词 {index}" for index in range(1, count + 1))


class PromptMultiFunctionTests(unittest.TestCase):
    def setUp(self):
        self.node = O1keyPromptMultiFunction()

    def test_schema_appends_multi_select_widgets_without_changing_legacy_positions(self):
        required = O1keyPromptMultiFunction.INPUT_TYPES()["required"]
        self.assertEqual(list(required), ["提示词", "功能", "抽取数量", "指定序号"])
        self.assertEqual(required["功能"][0], ["全部使用", "随机抽取n套", "指定序号"])
        self.assertEqual(required["抽取数量"][1]["default"], 3)
        self.assertEqual(required["指定序号"][1]["default"], "1,2,3")

    def test_legacy_random_one_mode_still_accepts_the_old_two_argument_call(self):
        with mock.patch.object(prompt_module.random, "choice", return_value="提示词 2"):
            result = self.node.process(_prompts(3), "随机抽取1套")
        self.assertEqual(result, ("提示词 2",))

    def test_random_many_samples_without_replacement_and_restores_source_order(self):
        with mock.patch.object(prompt_module.random, "sample", return_value=[4, 1, 3]) as sample:
            result = self.node.process(_prompts(5), "随机抽取n套", 3)
        sample.assert_called_once()
        population, count = sample.call_args.args
        self.assertEqual(list(population), [0, 1, 2, 3, 4])
        self.assertEqual(count, 3)
        self.assertEqual(result, ("提示词 2\n---\n提示词 4\n---\n提示词 5",))

    def test_random_many_rejects_a_count_larger_than_the_prompt_set(self):
        with self.assertRaisesRegex(ValueError, "抽取数量 4 超过当前提示词总数 3"):
            self.node.process(_prompts(3), "随机抽取n套", 4)

    def test_selected_indices_support_common_separators_ranges_and_input_order(self):
        result = self.node.process(
            _prompts(5),
            "指定序号",
            指定序号="5，2 3-4",
        )
        self.assertEqual(
            result,
            ("提示词 5\n---\n提示词 2\n---\n提示词 3\n---\n提示词 4",),
        )

    def test_selected_indices_reject_duplicates_and_out_of_range_values(self):
        with self.assertRaisesRegex(ValueError, "序号 2 重复"):
            self.node.process(_prompts(3), "指定序号", 指定序号="2,2")
        with self.assertRaisesRegex(ValueError, "序号 4 超出范围"):
            self.node.process(_prompts(3), "指定序号", 指定序号="1,4")

    def test_only_random_modes_disable_execution_caching(self):
        self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取1套")))
        self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取多套")))
        self.assertTrue(math.isnan(O1keyPromptMultiFunction.IS_CHANGED("a", "随机抽取n套")))
        selected_key = O1keyPromptMultiFunction.IS_CHANGED("a", "指定序号", 3, "1,3")
        self.assertEqual(selected_key, "指定序号|1,3|a")
        self.assertEqual(
            O1keyPromptMultiFunction.IS_CHANGED("a", "全部使用", 9, "2"),
            "全部使用|a",
        )


if __name__ == "__main__":
    unittest.main()