import pytest
from src.vendor_capabilities import VendorCapabilities, get_capabilities, register

@pytest.fixture(autouse=True)
def _clean_registry():
 import src.vendor_capabilities
 snapshot = src.vendor_capabilities._REGISTRY.copy()
 yield
 src.vendor_capabilities._REGISTRY.clear()
 src.vendor_capabilities._REGISTRY.update(snapshot)

def test_registry_lookup_known_model():
 caps = VendorCapabilities(
  vendor='qwen',
  model='qwen-max',
  vision=False,
  context_window=32768
 )
 register(caps)
 retrieved = get_capabilities('qwen', 'qwen-max')
 assert retrieved.vendor == 'qwen'
 assert retrieved.model == 'qwen-max'
 assert retrieved.context_window == 32768
 assert retrieved.vision is False

def test_fallback_to_vendor_default():
 caps = VendorCapabilities(
  vendor='llama',
  model='*',
  context_window=131072,
  cost_tracking=False
 )
 register(caps)
 retrieved = get_capabilities('llama', 'llama-3.3-future-unregistered')
 assert retrieved.context_window == 131072
 assert retrieved.cost_tracking is False

def test_unknown_vendor_raises():
 with pytest.raises(KeyError, match='No capabilities registered'):
  get_capabilities('nonexistent_vendor', 'anymodel')

V2_FIELDS: list[str] = [
 'local', 'reasoning', 'structured_output', 'code_execution',
 'web_search', 'x_search', 'file_search', 'mcp_support',
 'audio', 'video', 'grounding', 'computer_use',
]

@pytest.mark.parametrize('field_name', V2_FIELDS)
def test_v2_field_default_is_false(field_name: str) -> None:
 caps = VendorCapabilities(vendor='test', model='m')
 assert getattr(caps, field_name) is False, f'{field_name} should default to False'

@pytest.mark.parametrize('field_name', V2_FIELDS)
def test_v2_field_round_trip(field_name: str) -> None:
 caps = VendorCapabilities(vendor='test', model='m', **{field_name: True})
 assert getattr(caps, field_name) is True, f'{field_name} should round-trip to True'

def test_v2_local_flag_works_for_local_vendor() -> None:
 register(VendorCapabilities(vendor='llama', model='llama-local-test-3.1', local=True))
 caps = get_capabilities('llama', 'llama-local-test-3.1')
 assert caps.local is True

def test_v2_local_flag_falls_back_to_wildcard() -> None:
 register(VendorCapabilities(vendor='llama', model='*', local=True))
 caps = get_capabilities('llama', 'some-unregistered-model-3.1-future')
 assert caps.local is True

def test_v2_local_flag_does_not_affect_other_vendors() -> None:
 register(VendorCapabilities(vendor='llama', model='*', local=True))
 register(VendorCapabilities(vendor='qwen', model='*'))
 caps = get_capabilities('qwen', 'qwen-turbo')
 assert caps.local is False