import ast
import re
import unittest
from pathlib import Path


def load_helpers():
    source = Path(__file__).resolve().parents[1] / "ui_server.py"
    tree = ast.parse(source.read_text(encoding="utf-8"))
    names = {
        "_is_confirmation_prompt",
        "_is_ocean_export_request",
        "_apply_ocean_export_defaults",
    }
    functions = [
        node for node in tree.body
        if isinstance(node, ast.FunctionDef) and node.name in names
    ]
    namespace = {
        "re": re,
        "_needs_ocean_mcp": lambda prompt: any(
            term in str(prompt).lower()
            for term in ("era5", "ocean", "sst", "cmems", "oisst")
        ),
    }
    exec(
        compile(ast.Module(body=functions, type_ignores=[]), str(source), "exec"),
        namespace,
    )
    return namespace


helpers = load_helpers()


class FollowupRoutingTests(unittest.TestCase):
    def test_short_confirmation_variants(self):
        confirm = helpers["_is_confirmation_prompt"]
        for value in ("确认", "好的", "继续", "OK", "yes"):
            self.assertTrue(confirm(value))
        self.assertFalse(confirm("确认一下ERA5数据范围"))

    def test_missing_format_defaults_to_netcdf(self):
        apply_defaults = helpers["_apply_ocean_export_defaults"]
        prompt = "导出1998年1月4日ERA5 v10数据"
        routed = apply_defaults(prompt)
        self.assertIn("format=netcdf", routed)
        self.assertIn("mcp_marine_marine_export", routed)

    def test_explicit_format_is_preserved(self):
        apply_defaults = helpers["_apply_ocean_export_defaults"]
        prompt = "导出1998年1月4日ERA5 v10数据，格式csv"
        routed = apply_defaults(prompt)
        self.assertIn(prompt, routed)
        self.assertNotIn("format=netcdf", routed)
        self.assertIn("mcp_marine_marine_export", routed)


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