"""Bedrock cache provenance and compressor budgets across disk-backed restarts.

AWS documents Grok 4.6's Bedrock context window as 500K, independently of
xAI's direct API catalog: https://docs.aws.amazon.com/bedrock/latest/userguide/model-card-xai-grok-4-6.html
Only provider I/O is stubbed; cache/config readers and compressor are real.
"""

import json
import os
import subprocess
import sys
import time
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import yaml

from agent import bedrock_adapter as ba
from agent import model_metadata as mm
from agent.context_compressor import ContextCompressor


@pytest.fixture(autouse=True)
def isolated_home(tmp_path, monkeypatch):
    monkeypatch.setenv("HERMES_HOME", str(tmp_path))
    mm._BEDROCK_PROBE_FAILURE_CACHE.clear()
    monkeypatch.setattr(ba, "resolve_bedrock_region", lambda: "us-east-1")
    yield tmp_path
    mm._BEDROCK_PROBE_FAILURE_CACHE.clear()


@pytest.mark.parametrize("model", ["xai.grok-4.6", "global.xai.grok-4.6", "us.xai.grok-4.6"])
@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"])
@pytest.mark.parametrize("legacy", [None, 128_000, 700_000])
@pytest.mark.parametrize("probed", [None, 128_000, 800_000])
def test_bedrock_resolution_migrates_ambiguous_cache_and_preserves_probe(
    isolated_home, monkeypatch, model, base_url, legacy, probed,
):
    """Neither an old small nor large scalar proves where it came from.

    A successful probe is authoritative, including below-table limits. Its
    provenance must survive unrelated writes and an actual process restart.
    """
    cache_url = base_url or "bedrock://"
    key = mm._context_cache_key(model, cache_url)
    cache_file = isolated_home / "context_length_cache.yaml"
    if legacy is not None:
        cache_file.write_text(yaml.safe_dump({"context_lengths": {key: legacy}}))
    probe = Mock(return_value=probed)
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    expected = probed if probed is not None else 500_000
    compressor = ContextCompressor(model, provider="bedrock", base_url=base_url, quiet_mode=True)
    assert compressor.context_length == expected
    # Preserve the existing raise-only 75% floor for windows below 512K.
    expected_threshold = int(expected * (0.75 if expected < 512_000 else 0.5))
    assert compressor.threshold_tokens == expected_threshold
    assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == expected
    probe.assert_called_once_with(model, "us-east-1")
    assert mm.get_cached_context_length(model, cache_url, bedrock_confirmed=True) == probed
    # Ordinary cache updates must not erase another entry's provenance.
    mm.save_context_length("other-model", "https://other.example/v1", 64_000)
    mm._invalidate_cached_context_length("other-model", "https://other.example/v1")
    if probed is None:
        assert yaml.safe_load(cache_file.read_text())["context_lengths"].get(key) == legacy
    else:
        script = '''
import json, sys
from agent import bedrock_adapter as ba
from agent.context_compressor import ContextCompressor
ba.probe_bedrock_context_length = lambda *a, **k: (_ for _ in ()).throw(AssertionError("reprobed persisted success"))
c = ContextCompressor(sys.argv[1], provider="bedrock", base_url=sys.argv[2], quiet_mode=True)
print(json.dumps([c.context_length, c.threshold_tokens]))
'''
        result = subprocess.run([sys.executable, "-c", script, model, base_url],
                                env=dict(os.environ), text=True, capture_output=True, check=True)
        assert json.loads(result.stdout) == [expected, expected_threshold]


@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"])
@pytest.mark.parametrize("retry", ["ttl", "invalidate", "restart", "profile", "endpoint"])
def test_failure_memo_retry_scope_and_expiry(isolated_home, monkeypatch, base_url, retry):
    model = "global.xai.grok-4.6"
    cache_url = base_url or "bedrock://"
    probe = Mock(side_effect=[None, 96_000])
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000
    assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000
    assert probe.call_count == 1
    assert not (isolated_home / "context_length_cache.yaml").exists()
    if retry == "ttl":
        expired = time.monotonic() - mm._BEDROCK_PROBE_FAILURE_TTL_SECONDS - 1
        for key in mm._BEDROCK_PROBE_FAILURE_CACHE:
            mm._BEDROCK_PROBE_FAILURE_CACHE[key] = expired
    elif retry == "invalidate":
        mm._invalidate_cached_context_length(model, cache_url)
    elif retry == "restart":
        mm._BEDROCK_PROBE_FAILURE_CACHE.clear()
    elif retry == "profile":
        new_home = isolated_home / "other-profile"
        new_home.mkdir()
        monkeypatch.setenv("HERMES_HOME", str(new_home))
    else:
        base_url = "https://bedrock-runtime.us-east-1.amazonaws.com/other"
    assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 96_000
    assert probe.call_count == 2
    if retry == "ttl":
        assert not mm._BEDROCK_PROBE_FAILURE_CACHE


@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"])
@pytest.mark.parametrize("override", ["argument", "config"])
def test_explicit_context_override_and_compressor_caps_win(isolated_home, monkeypatch, base_url, override):
    model = "us.xai.grok-4.6"
    probe = Mock(side_effect=AssertionError("explicit override must not probe"))
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    explicit_context = 80_000 if override == "argument" else None
    if override == "config":
        (isolated_home / "config.yaml").write_text(yaml.safe_dump({
            "model_overrides": {"bedrock": {model: {"context_window": 80_000}}},
        }))
    compressor = ContextCompressor(model, provider="bedrock", base_url=base_url,
                                   threshold_tokens_cap=30_000, max_tokens=10_000,
                                   quiet_mode=True, config_context_length=explicit_context)
    assert compressor.context_length == 80_000
    assert compressor.threshold_tokens == 30_000
    assert not (isolated_home / "context_length_cache.yaml").exists()
    probe.assert_not_called()


def test_unknown_model_fallback_and_host_inference(monkeypatch):
    probe = Mock(return_value=None)
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    base_url = "https://bedrock-runtime.us-east-1.amazonaws.com"
    assert mm.get_model_context_length("unknown.future-model", base_url=base_url) == ba.BEDROCK_DEFAULT_CONTEXT_LENGTH
    assert mm.get_model_context_length("xai.grok-4.6", base_url=base_url) == 500_000
    assert mm.get_cached_context_length("unknown.future-model", base_url) is None


@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"])
@pytest.mark.parametrize("writer", ["overflow", "usage"])
def test_provider_confirmed_writers_survive_restart(monkeypatch, base_url, writer):
    from agent.turn_overflow import _adopt_provider_context_limit
    from agent.turn_usage import record_response_usage

    model = "global.xai.grok-4.6"
    compressor = ContextCompressor(model, base_url=base_url, provider="bedrock", quiet_mode=True)
    compressor.context_length = 500_000
    agent = SimpleNamespace(model=model, provider="bedrock", api_mode="bedrock", base_url=base_url,
                            context_compressor=compressor, _buffer_vprint=lambda *a: None,
                            _safe_print=lambda *a: None, log_prefix="", client=None,
                            _session_db=None, verbose_logging=False, quiet_mode=True,
                            session_api_calls=0, session_estimated_cost_usd=0)
    if writer == "overflow":
        assert _adopt_provider_context_limit(SimpleNamespace(agent=agent),
            "maximum context length is 96000 tokens", 500_000) == 96_000
    else:
        compressor.context_length = 96_000
        compressor._context_probed = compressor._context_probe_persistable = True
        for name in ("prompt", "completion", "total", "input", "output", "cache_read", "cache_write", "reasoning"):
            setattr(agent, f"session_{name}_tokens", 0)
        response = SimpleNamespace(usage={"input_tokens": 100, "output_tokens": 5})
        record_response_usage(agent, response, messages=[{"role": "user", "content": "hi"}],
                              api_call_count=1, api_duration=0.1, compression_attempts=0, max_compression_attempts=3)
    assert mm.get_cached_context_length(model, base_url or "bedrock://") == 96_000
    script = '''
import sys
from agent import model_metadata as mm, bedrock_adapter as ba
ba.probe_bedrock_context_length = lambda *a, **k: (_ for _ in ()).throw(AssertionError("lost provider limit"))
assert mm.get_model_context_length(sys.argv[1], base_url=sys.argv[2], provider="bedrock") == 96000
'''
    subprocess.run([sys.executable, "-c", script, model, base_url], check=True, capture_output=True, text=True)


@pytest.mark.parametrize("base_url", ["", "https://bedrock-runtime.us-east-1.amazonaws.com"])
def test_readonly_legacy_cache_does_not_reset_probe_cooldown(isolated_home, monkeypatch, base_url):
    model = "global.xai.grok-4.6"
    key = mm._context_cache_key(model, base_url or "bedrock://")
    cache_file = isolated_home / "context_length_cache.yaml"
    cache_file.write_text(yaml.safe_dump({"context_lengths": {key: 128_000}}))
    probe = Mock(return_value=None)
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    monkeypatch.setattr(mm, "_write_context_cache", Mock(side_effect=OSError("read-only")))
    for _ in range(3):
        assert mm.get_model_context_length(model, provider="bedrock", base_url=base_url) == 500_000
    assert probe.call_count == 1


@pytest.mark.parametrize("rewrite", [None, 96_000, 120_000])
def test_provenance_is_backward_readable_and_generic_writes_clear_it(isolated_home, monkeypatch, rewrite):
    model, base_url = "xai.grok-4.6", "https://bedrock-runtime.us-east-1.amazonaws.com"
    probe = Mock(side_effect=[96_000, 110_000])
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    assert mm.get_model_context_length(model, base_url=base_url) == 96_000
    raw = yaml.safe_load((isolated_home / "context_length_cache.yaml").read_text())
    # Old readers take this value directly into arithmetic. Metadata is additive.
    assert raw["context_lengths"][mm._context_cache_key(model, base_url)] + 1 == 96_001
    if rewrite is not None:
        mm.save_context_length(model, base_url, rewrite)
        assert mm.get_model_context_length(model, base_url=base_url) == 110_000
    else:
        assert mm.get_model_context_length(model, base_url=base_url) == 96_000


@pytest.mark.parametrize("lengths", [True, [128_000], "bad"])
def test_malformed_lengths_do_not_block_provider_persistence(isolated_home, monkeypatch, lengths):
    (isolated_home / "context_length_cache.yaml").write_text(yaml.safe_dump({"context_lengths": lengths}))
    monkeypatch.setattr(ba, "probe_bedrock_context_length", lambda *a: 96_000)
    assert mm.get_model_context_length("xai.grok-4.6", provider="bedrock") == 96_000
    assert mm.get_cached_context_length("xai.grok-4.6", "bedrock://") == 96_000


@pytest.mark.parametrize("marker", [True, "96000", 128_000, [96_000], {"source": "probe"}])
def test_malformed_or_mismatched_provenance_requires_revalidation(isolated_home, monkeypatch, marker):
    model, base_url = "xai.grok-4.6", "bedrock://"
    key = mm._context_cache_key(model, base_url)
    (isolated_home / "context_length_cache.yaml").write_text(yaml.safe_dump({
        "context_lengths": {key: 96_000}, "bedrock_confirmed_v1": {key: marker},
    }))
    probe = Mock(return_value=100_000)
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    assert mm.get_model_context_length(model, provider="bedrock") == 100_000
    probe.assert_called_once()


def test_context_local_profile_memos_do_not_cross_and_expired_rows_are_pruned(isolated_home, monkeypatch):
    from hermes_constants import set_hermes_home_override, reset_hermes_home_override

    probe = Mock(side_effect=[None, None, 96_000])
    monkeypatch.setattr(ba, "probe_bedrock_context_length", probe)
    model = "global.xai.grok-4.6"
    assert mm.get_model_context_length(model, provider="bedrock") == 500_000
    # Leave a different model's expired row: lookup must prune it too.
    assert mm.get_model_context_length("unknown.future", provider="bedrock") == ba.BEDROCK_DEFAULT_CONTEXT_LENGTH
    for key in mm._BEDROCK_PROBE_FAILURE_CACHE:
        if "unknown.future" in key:
            mm._BEDROCK_PROBE_FAILURE_CACHE[key] = time.monotonic() - mm._BEDROCK_PROBE_FAILURE_TTL_SECONDS - 1
    token = set_hermes_home_override(isolated_home / "routed-profile")
    try:
        assert mm.get_model_context_length(model, provider="bedrock") == 96_000
    finally:
        reset_hermes_home_override(token)
    assert probe.call_count == 3
    assert all("unknown.future" not in key for key in mm._BEDROCK_PROBE_FAILURE_CACHE)
    assert mm.get_model_context_length(model, provider="bedrock") == 500_000
    assert probe.call_count == 3