import asyncio
import unittest
from unittest.mock import AsyncMock, patch
import app
from gradio.mcp import GradioMCPServer


class AdapterTests(unittest.IsolatedAsyncioTestCase):
    async def test_only_six_read_tools_exposed(self):
        server = GradioMCPServer(app.demo)
        self.assertEqual(set(server.tool_to_endpoint), app.ALLOWED)
        self.assertEqual(len(app.demo.fns), 6)
        for fn in app.demo.fns.values():
            self.assertEqual(fn.concurrency_limit, 4)
            self.assertEqual(fn.concurrency_id, 'public_reads')

    async def test_unsafe_tool_rejected_before_network(self):
        with patch.object(app.httpx, 'AsyncClient') as client:
            with self.assertRaisesRegex(ValueError, 'not allowed'):
                await app.forward('run', {'dataset': 'lerobot/pusht'})
            client.assert_not_called()

    async def test_bounds_and_unknown_parameters_rejected_before_network(self):
        for name,args in [('list_datasets',{'limit':1000}),('get_report',{'url':'http://localhost'}),
                          ('get_findings',{'dataset':'lerobot/pusht'}),
                          ('search_datasets',{'q':'a'*129}),('search_episodes',{'limit':31}),('search_combined_datasets',{'per_page':1}),('search_combined_datasets',{'direction':'bad'})]:
            with self.subTest(name=name),patch.object(app.httpx, 'AsyncClient') as client:
                with self.assertRaisesRegex(ValueError, 'Invalid tool arguments'):
                    await app.forward(name,args)
                client.assert_not_called()

    async def test_wrappers_preserve_revision_and_pagination(self):
        result={'content':[], 'structuredContent':{'report_revision':'r1'},'isError':False}
        with patch.object(app,'forward',new=AsyncMock(return_value=result)) as forward:
            self.assertEqual(await app.get_findings('r1',dataset='lerobot/pusht',after=5,limit=3),result)
            name,args=forward.call_args.args
            self.assertEqual(name,'get_findings')
            self.assertEqual(args['report_revision'],'r1')
            self.assertEqual(args['after'],5)
            self.assertEqual(args['limit'],3)

    async def test_fixed_anonymous_origin_and_no_environment_auth(self):
        with patch.object(app.httpx,'AsyncClient',side_effect=TimeoutError) as client:
            with self.assertRaisesRegex(app.gr.Error,'No audit was started'):
                await app.get_report(dataset='lerobot/pusht')
            self.assertFalse(client.call_args.kwargs['trust_env'])
            self.assertFalse(client.call_args.kwargs['follow_redirects'])
            self.assertNotIn('headers',client.call_args.kwargs)
        self.assertEqual(app.UPSTREAM,'https://mcp.kalanos.ai/mcp')

    async def test_transport_preserves_json_and_upstream_errors(self):
        import json
        from contextlib import asynccontextmanager
        from unittest.mock import Mock
        @asynccontextmanager
        async def transport(*args, **kwargs):
            yield None, None, None
        @asynccontextmanager
        async def session_context(*args):
            yield session
        session=Mock(initialize=AsyncMock(),call_tool=AsyncMock())
        for failed in (False,True):
            envelope={'content':[], 'structuredContent':{'report_revision':'r1'},'isError':failed}
            result=Mock();result.model_dump.return_value=envelope
            session.call_tool.return_value=result
            with patch.object(app,'streamable_http_client',transport),patch.object(app,'ClientSession',session_context):
                if failed:
                    with self.assertRaisesRegex(app.gr.Error,'Upstream read-only error'):
                        await app.get_report(report_id='report1')
                else:
                    self.assertEqual(json.loads(await app.get_report(report_id='report1')),envelope)
            session.call_tool.assert_awaited_with('get_report',{'report_id':'report1'})

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


class CombinedAdapterTests(unittest.IsolatedAsyncioTestCase):
    async def test_combined_wrapper_forwards_filter_and_paging(self):
        with patch.object(app, 'forward', new=AsyncMock(return_value='ok')) as f:
            self.assertEqual(await app.search_combined_datasets(task=['pick'],robot=['aloha'],page=2,per_page=40,direction='asc',facets=True),'ok')
            name,args=f.call_args.args
            self.assertEqual(name,'search_combined_datasets')
            self.assertEqual([args[k] for k in ['task','robot','page','per_page','direction','facets']],[['pick'],['aloha'],2,40,'asc',True])
    async def test_schema_matches_checked_upstream_contract(self):
        import json
        from pathlib import Path
        from jsonschema import Draft7Validator
        schema=app.SCHEMAS['search_combined_datasets'];Draft7Validator.check_schema(schema)
        self.assertFalse(list(Draft7Validator(schema).iter_errors({'task':['pick'],'page':2,'per_page':40,'direction':'asc'})))