import unittest
from pathlib import Path
import yaml
from render_config import build_config


class PolicyTests(unittest.TestCase):
    def build(self, **env):
        catalog = yaml.safe_load(Path(__file__).with_name('models.yaml').read_text())
        return build_config(catalog, env, lambda _: True)

    def test_subscription_is_ignored_even_if_old_env_still_exists(self):
        c = self.build(ZENMUX_SUBSCRIPTION_API_KEY='retired-test')
        self.assertNotIn('ZENMUX_SUBSCRIPTION_API_KEY', yaml.safe_dump(c))
        self.assertFalse(any('zenmux.ai' in r['litellm_params'].get('api_base', '') for r in c['model_list']))

    def test_openrouter_primary_and_byok_last_for_all_supported_aliases(self):
        c = self.build(ZENMUX_SUBSCRIPTION_API_KEY='retired', ZENMUX_PAYG_API_KEY='test', OPENROUTER_API_KEY='test')
        routes = {r['model_name']: r['litellm_params'] for r in c['model_list']}
        groups = {k:v for row in c['router_settings']['fallbacks'] for k,v in row.items()}
        for alias in ['claude-sonnet-latest', 'claude-opus-latest', 'claude-haiku-latest',
                      'deepseek-flash-latest', 'glm-latest', 'kimi-latest', 'mimo-flash-latest']:
            for name in [alias, 'auto/'+alias, 'zenmux/'+alias, 'openrouter/'+alias]:
                self.assertEqual(routes[name]['api_key'], 'os.environ/OPENROUTER_API_KEY')
                self.assertEqual(groups[name][-1], 'zenmux-payg/'+alias)
            self.assertEqual(routes['zenmux-payg/'+alias]['api_key'], 'os.environ/ZENMUX_PAYG_API_KEY')
        for source, targets in groups.items():
            self.assertTrue(all(target in routes for target in targets))
            self.assertNotEqual(routes[source].get('api_key'), 'os.environ/ZENMUX_PAYG_API_KEY')
        self.assertEqual(len(routes), len(c['model_list']))

    def test_no_byok_key_no_dangling_fallback_and_no_subscription(self):
        c = self.build(OPENROUTER_API_KEY='test', ZENMUX_SUBSCRIPTION_API_KEY='retired')
        self.assertEqual(c['router_settings']['fallbacks'], [])
        self.assertNotIn('ZENMUX_SUBSCRIPTION_API_KEY', yaml.safe_dump(c))

    def test_oauth_official_fallback_stays_unchanged(self):
        c = self.build(OPENAI_API_KEY='test', OPENROUTER_API_KEY='test', ZENMUX_PAYG_API_KEY='test')
        groups = {k:v for row in c['router_settings']['fallbacks'] for k,v in row.items()}
        self.assertEqual(groups['codex/gpt-sol-latest'], ['openai/gpt-sol-latest'])

    def test_google_primary_before_gemini_byok(self):
        c = self.build(GOOGLE_AI_STUDIO_API_KEY='test', ZENMUX_PAYG_API_KEY='test')
        routes = {r['model_name']: r['litellm_params'] for r in c['model_list']}
        self.assertEqual(routes['zenmux/gemini-flash-latest']['api_key'], 'os.environ/GOOGLE_AI_STUDIO_API_KEY')


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