import asyncio
import json
import unittest
from typing import Optional

import httpx
from pwo import async_test

from kaya.core import HttpContext, KayaApp
from kaya.core._asgi import AsgiWebSocket
from kaya.forwarded import ForwardedHeadersMixin

TRUSTED = ['127.0.0.1', '::1', '10.0.0.0/8']


def make_app(trusted_proxies=TRUSTED) -> KayaApp:
    mixins = [ForwardedHeadersMixin(trusted_proxies=trusted_proxies)] if trusted_proxies is not None else []
    app = KayaApp(mixins=mixins)

    @app.GET('/client')
    async def client(ctx: HttpContext) -> None:
        host, port = ctx.client if ctx.client is not None else (None, None)
        await ctx.send_str(200, json.dumps({'host': host, 'port': port}))

    return app


async def request(app: KayaApp,
                  headers: Optional[dict[str, str]] = None,
                  client: tuple[str, int] = ('127.0.0.1', 123)) -> dict:
    transport = httpx.ASGITransport(app=app, client=client)
    async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1:80") as http_client:
        r = await http_client.get('/client', headers=headers)
        assert r.status_code == 200
        return json.loads(r.text)


class TrustedPeerTest(unittest.TestCase):

    @async_test
    async def test_no_headers_returns_socket_peer(self):
        app = make_app()
        result = await request(app)
        self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)

    @async_test
    async def test_forwarded_header_with_port(self):
        app = make_app()
        result = await request(app, headers={'Forwarded': 'for=203.0.113.5:1234'})
        self.assertEqual({'host': '203.0.113.5', 'port': 1234}, result)

    @async_test
    async def test_forwarded_header_bracketed_ipv6(self):
        app = make_app()
        result = await request(app, headers={'Forwarded': 'for="[2001:db8::1]:4711"'})
        self.assertEqual({'host': '2001:db8::1', 'port': 4711}, result)

    @async_test
    async def test_forwarded_header_without_port_keeps_socket_port(self):
        app = make_app()
        result = await request(app, headers={'Forwarded': 'for=203.0.113.5'})
        self.assertEqual({'host': '203.0.113.5', 'port': 123}, result)

    @async_test
    async def test_forwarded_unknown_entry_skipped(self):
        app = make_app()
        result = await request(app, headers={'Forwarded': 'for=unknown, for=203.0.113.5'})
        self.assertEqual('203.0.113.5', result['host'])

    @async_test
    async def test_forwarded_rightmost_untrusted_wins(self):
        # attacker-controlled leftmost entry is skipped: the rightmost
        # untrusted entry (appended by the trusted edge proxy) is the client
        app = make_app()
        result = await request(app, headers={'Forwarded': 'for=1.2.3.4, for=5.6.7.8, for=10.0.0.2'})
        self.assertEqual('5.6.7.8', result['host'])

    @async_test
    async def test_x_forwarded_for_spoofed_leftmost_entry_skipped(self):
        # XFF = "<attacker-supplied>, <real client>" as appended by the proxy
        app = make_app()
        result = await request(app, headers={'X-Forwarded-For': '1.2.3.4, 5.6.7.8'})
        self.assertEqual('5.6.7.8', result['host'])

    @async_test
    async def test_x_forwarded_for_all_trusted_chain_uses_leftmost(self):
        app = make_app()
        result = await request(app, headers={'X-Forwarded-For': '10.0.0.5, 10.0.0.2'})
        self.assertEqual('10.0.0.5', result['host'])

    @async_test
    async def test_x_forwarded_port(self):
        app = make_app()
        result = await request(app, headers={
            'X-Forwarded-For': '203.0.113.5',
            'X-Forwarded-Port': '8443',
        })
        self.assertEqual({'host': '203.0.113.5', 'port': 8443}, result)

    @async_test
    async def test_invalid_x_forwarded_port_ignored(self):
        app = make_app()
        result = await request(app, headers={
            'X-Forwarded-For': '203.0.113.5',
            'X-Forwarded-Port': 'not-a-port',
        })
        self.assertEqual({'host': '203.0.113.5', 'port': 123}, result)

    @async_test
    async def test_forwarded_takes_precedence_over_x_forwarded_for(self):
        app = make_app()
        result = await request(app, headers={
            'Forwarded': 'for=203.0.113.5',
            'X-Forwarded-For': '198.51.100.7',
        })
        self.assertEqual('203.0.113.5', result['host'])

    @async_test
    async def test_x_forwarded_host_fallback(self):
        app = make_app()
        result = await request(app, headers={'X-Forwarded-Host': '198.51.100.7'})
        self.assertEqual('198.51.100.7', result['host'])

    @async_test
    async def test_ipv6_cidr_trust(self):
        app = make_app(trusted_proxies=['2001:db8::/32'])
        result = await request(app,
                               headers={'X-Forwarded-For': '203.0.113.5'},
                               client=('2001:db8::10', 9999))
        self.assertEqual({'host': '203.0.113.5', 'port': 9999}, result)


class UntrustedPeerTest(unittest.TestCase):

    @async_test
    async def test_untrusted_peer_ignores_forwarded_headers(self):
        app = make_app()
        result = await request(app,
                               headers={'Forwarded': 'for=1.2.3.4', 'X-Forwarded-For': '1.2.3.4'},
                               client=('203.0.113.99', 4567))
        self.assertEqual({'host': '203.0.113.99', 'port': 4567}, result)

    @async_test
    async def test_empty_trusted_proxies_ignores_everything(self):
        app = make_app(trusted_proxies=[])
        result = await request(app, headers={'X-Forwarded-For': '1.2.3.4'})
        self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)

    @async_test
    async def test_invalid_cidr_fails_fast(self):
        with self.assertRaises(ValueError):
            ForwardedHeadersMixin(trusted_proxies=['not-a-cidr'])


class OptOutTest(unittest.TestCase):

    @async_test
    async def test_without_mixin_headers_are_ignored(self):
        app = KayaApp()

        @app.GET('/client')
        async def client(ctx: HttpContext) -> None:
            host, port = ctx.client if ctx.client is not None else (None, None)
            await ctx.send_str(200, json.dumps({'host': host, 'port': port}))

        result = await request(app, headers={'Forwarded': 'for=1.2.3.4', 'X-Forwarded-For': '1.2.3.4'})
        self.assertEqual({'host': '127.0.0.1', 'port': 123}, result)


class WebSocketTest(unittest.TestCase):

    @staticmethod
    def _make_ws(headers):
        async def send(message):
            pass

        async def receive():
            return {'type': 'websocket.connect'}

        scope = {
            'type': 'websocket',
            'path': '/ws',
            'query_string': b'',
            'scheme': 'ws',
            'client': ('127.0.0.1', 12345),
            'server': ('127.0.0.1', 80),
            'headers': headers,
        }
        return AsgiWebSocket(scope, receive, send)

    @async_test
    async def test_websocket_trusted_peer(self):
        mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
        ws = self._make_ws([(b'x-forwarded-for', b'1.2.3.4, 5.6.7.8')])
        wrapped = await mixin._before_websocket(ws)
        assert wrapped is not None
        self.assertEqual(('5.6.7.8', 12345), wrapped.client)

    @async_test
    async def test_websocket_untrusted_peer(self):
        mixin = ForwardedHeadersMixin(trusted_proxies=['10.0.0.0/8'])
        ws = self._make_ws([(b'x-forwarded-for', b'1.2.3.4')])
        wrapped = await mixin._before_websocket(ws)
        self.assertIsNone(wrapped)
        self.assertEqual(('127.0.0.1', 12345), ws.client)


class RsgiTest(unittest.TestCase):

    def test_rsgi_context(self):
        from kaya.rsgi import RsgiContext

        class FakeScope:
            scheme = 'http'
            method = 'GET'
            path = '/'
            query_string = ''
            headers = {'x-forwarded-for': '1.2.3.4, 5.6.7.8'}
            client = '127.0.0.1:12345'
            server = '127.0.0.1:80'

        mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
        ctx = RsgiContext(FakeScope(), object())  # type: ignore[arg-type]
        wrapped = asyncio.run(mixin._before_request(ctx))
        assert wrapped is not None
        self.assertEqual(('5.6.7.8', 12345), wrapped.client)

    def test_rsgi_context_untrusted_peer(self):
        from kaya.rsgi import RsgiContext

        class FakeScope:
            scheme = 'http'
            method = 'GET'
            path = '/'
            query_string = ''
            headers = {'x-forwarded-for': '1.2.3.4'}
            client = '192.0.2.1:12345'
            server = '127.0.0.1:80'

        mixin = ForwardedHeadersMixin(trusted_proxies=TRUSTED)
        ctx = RsgiContext(FakeScope(), object())  # type: ignore[arg-type]
        wrapped = asyncio.run(mixin._before_request(ctx))
        self.assertIsNone(wrapped)
        self.assertEqual(('192.0.2.1', 12345), ctx.client)