from concurrent.futures import ThreadPoolExecutor
from types import SimpleNamespace
from uuid import uuid4
import re

from artemis.llm.cache_affinity import cache_affinity_options


def context(session, exploration=None):
    return SimpleNamespace(
        data_engine=SimpleNamespace(current_session_id=session), exploration_id=exploration
    )


def test_routing_identity_survives_resume_but_isolates_tasks_roles_and_models():
    ctx = context(uuid4(), "exploration")
    first = cache_affinity_options(ctx, "flash", "gpt-6-luna")
    ctx.data_engine.current_session_id = uuid4()
    assert first == cache_affinity_options(ctx, "flash", "gpt-6-luna")
    key = first["prompt_cache_key"]
    assert re.fullmatch(r"[-0-9A-Za-z._:]{1,128}", key)
    assert "exploration" not in key
    assert first["extra_headers"] == {"Thread-Id": key, "Session-Id": key}
    for role, model in [("checker", "gpt-6-luna"), ("flash", "gpt-6.1-sol")]:
        assert cache_affinity_options(ctx, role, model)["prompt_cache_key"] != key
    assert (
        cache_affinity_options(context("another"), "flash", "gpt-6-luna")["prompt_cache_key"] != key
    )


def test_explicit_options_are_preserved_without_mutating_headers_or_adding_aliases():
    headers = {"thread_ID": "explicit-thread", "session-id": "explicit-session", "X-Test": "ok"}
    options = {"prompt_cache_key": "explicit.key", "extra_headers": headers, "temperature": 0}
    actual = cache_affinity_options(context("task"), "flash", "test", options)
    assert actual == options and actual is not options
    assert actual["extra_headers"] is not headers
    assert options["extra_headers"] == headers
    assert cache_affinity_options(None, "flash", "test") == {}
    for key in (None, "", "a" * 129, "unsafe\r\nheader"):
        assert cache_affinity_options(
            context("task"), "flash", "test", {"prompt_cache_key": key}
        ) == {"prompt_cache_key": key}


def test_concurrent_calls_do_not_share_or_mutate_task_headers():
    shared = {"extra_headers": {"X-Test": "shared"}}
    with ThreadPoolExecutor(max_workers=2) as pool:
        results = list(
            pool.map(
                lambda task: cache_affinity_options(context(task), "flash", "test", shared),
                ["task-a", "task-b"],
            )
        )
    assert results[0]["extra_headers"]["Thread-Id"] != results[1]["extra_headers"]["Thread-Id"]
    assert shared == {"extra_headers": {"X-Test": "shared"}}