"""Proxy and Outpost e2e tests"""

from base64 import b64encode
from dataclasses import asdict
from json import dumps
from time import sleep

from jwt import decode
from selenium.webdriver.common.by import By

from authentik.blueprints.tests import apply_blueprint, reconcile_app
from authentik.core.models import Application
from authentik.flows.models import Flow
from authentik.lib.generators import generate_id
from authentik.outposts.models import Outpost, OutpostConfig, OutpostType
from authentik.providers.proxy.models import ProxyProvider
from tests.decorators import retry
from tests.live import ChannelsE2ETestCase
from tests.selenium import SeleniumTestCase


class TestProviderProxy(SeleniumTestCase):
    """Proxy and Outpost e2e tests"""

    def setUp(self):
        super().setUp()
        self.run_container(
            image="traefik/whoami:latest",
            ports={
                "80": "80",
            },
        )

    def start_proxy(self, outpost: Outpost):
        """Start proxy container based on outpost created"""
        self.run_container(
            image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
            ports={
                "9000": "9000",
            },
            environment={
                "AUTHENTIK_TOKEN": outpost.token.key,
            },
        )

    @retry()
    @apply_blueprint(
        "default/flow-default-authentication-flow.yaml",
        "default/flow-default-invalidation-flow.yaml",
    )
    @apply_blueprint(
        "default/flow-default-provider-authorization-implicit-consent.yaml",
        "default/flow-default-provider-invalidation.yaml",
    )
    @apply_blueprint(
        "system/providers-oauth2.yaml",
        "system/providers-proxy.yaml",
    )
    @reconcile_app("authentik_crypto")
    def test_proxy_simple(self):
        """Test simple outpost setup with single provider"""
        # set additionalHeaders to test later
        self.user.attributes["additionalHeaders"] = {"X-Foo": "bar"}
        self.user.save()

        proxy: ProxyProvider = ProxyProvider.objects.create(
            name=generate_id(),
            authorization_flow=Flow.objects.get(
                slug="default-provider-authorization-implicit-consent"
            ),
            invalidation_flow=Flow.objects.get(slug="default-provider-invalidation-flow"),
            internal_host=f"http://{self.host}",
            external_host="http://localhost:9000",
        )
        # Ensure OAuth2 Params are set
        proxy.set_oauth_defaults()
        proxy.save()
        # we need to create an application to actually access the proxy
        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)
        outpost: Outpost = Outpost.objects.create(
            name=generate_id(),
            type=OutpostType.PROXY,
        )
        outpost.providers.add(proxy)
        outpost.build_user_permissions(outpost.user)

        self.start_proxy(outpost)

        sleep(5)

        self.driver.get("http://localhost:9000/api")
        self.login()
        sleep(1)

        body = self.parse_json_content()
        headers = body.get("headers", {})
        snippet = dumps(body, indent=2)[:500].replace("\n", " ")

        self.assertEqual(
            headers.get("X-Authentik-Username"),
            [self.user.username],
            f"X-Authentik-Username header mismatch at {self.driver.current_url}: {snippet}",
        )
        self.assertEqual(
            headers.get("X-Foo"),
            ["bar"],
            f"X-Foo header mismatch at {self.driver.current_url}: {snippet}",
        )
        raw_jwt: str = headers.get("X-Authentik-Jwt", [None])[0]
        jwt = decode(raw_jwt, options={"verify_signature": False})

        self.assertIsNotNone(jwt["sid"], "Missing 'sid' in JWT")
        self.assertIsNotNone(jwt["ak_proxy"], "Missing 'ak_proxy' in JWT")

        self.driver.get("http://localhost:9000/outpost.goauthentik.io/sign_out")
        sleep(2)

        flow_executor = self.get_shadow_root("ak-flow-executor")
        session_end_stage = self.get_shadow_root("ak-stage-session-end", flow_executor)
        flow_card = self.get_shadow_root("ak-flow-card", session_end_stage)
        title = flow_card.find_element(By.CSS_SELECTOR, ".pf-c-title.pf-m-3xl").text

        self.assertIn(
            "You've logged out of",
            title,
            f"Logout title mismatch at {self.driver.current_url}: {title}",
        )

    @retry()
    @apply_blueprint(
        "default/flow-default-authentication-flow.yaml",
        "default/flow-default-invalidation-flow.yaml",
    )
    @apply_blueprint(
        "default/flow-default-provider-authorization-implicit-consent.yaml",
        "default/flow-default-provider-invalidation.yaml",
    )
    @apply_blueprint(
        "system/providers-oauth2.yaml",
        "system/providers-proxy.yaml",
    )
    @reconcile_app("authentik_crypto")
    def test_proxy_basic_auth(self):
        """Test simple outpost setup with single provider"""
        cred = generate_id()
        attr = "basic-password"  # nosec
        self.user.attributes["basic-username"] = cred
        self.user.attributes[attr] = cred
        self.user.save()

        proxy: ProxyProvider = ProxyProvider.objects.create(
            name=generate_id(),
            authorization_flow=Flow.objects.get(
                slug="default-provider-authorization-implicit-consent"
            ),
            invalidation_flow=Flow.objects.get(slug="default-provider-invalidation-flow"),
            internal_host=f"http://{self.host}",
            external_host="http://localhost:9000",
            basic_auth_enabled=True,
            basic_auth_user_attribute="basic-username",
            basic_auth_password_attribute=attr,
        )
        # Ensure OAuth2 Params are set
        proxy.set_oauth_defaults()
        proxy.save()
        # we need to create an application to actually access the proxy
        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)
        outpost: Outpost = Outpost.objects.create(
            name=generate_id(),
            type=OutpostType.PROXY,
        )
        outpost.providers.add(proxy)
        outpost.build_user_permissions(outpost.user)

        self.start_proxy(outpost)

        sleep(5)

        self.driver.get("http://localhost:9000/api")
        self.login()
        sleep(1)

        body = self.parse_json_content()
        headers = body.get("headers", {})
        snippet = dumps(body, indent=2)[:500].replace("\n", " ")

        self.assertEqual(
            headers.get("X-Authentik-Username"),
            [self.user.username],
            f"X-Authentik-Username header mismatch at {self.driver.current_url}: {snippet}",
        )

        auth_header = b64encode(f"{cred}:{cred}".encode()).decode()

        self.assertEqual(
            headers.get("Authorization"),
            [f"Basic {auth_header}"],
            f"Authorization header mismatch at {self.driver.current_url}: {snippet}",
        )

        self.driver.get("http://localhost:9000/outpost.goauthentik.io/sign_out")
        sleep(2)

        flow_executor = self.get_shadow_root("ak-flow-executor")
        session_end_stage = self.get_shadow_root("ak-stage-session-end", flow_executor)
        flow_card = self.get_shadow_root("ak-flow-card", session_end_stage)
        title = flow_card.find_element(By.CSS_SELECTOR, ".pf-c-title.pf-m-3xl").text

        self.assertIn(
            "You've logged out of",
            title,
            f"Logout title mismatch at {self.driver.current_url}: {title}",
        )


class TestProviderProxyConnect(ChannelsE2ETestCase):
    """Test Proxy connectivity over websockets"""

    @retry(exceptions=[AssertionError])
    @apply_blueprint(
        "default/flow-default-authentication-flow.yaml",
        "default/flow-default-invalidation-flow.yaml",
    )
    @apply_blueprint(
        "default/flow-default-provider-authorization-implicit-consent.yaml",
    )
    @reconcile_app("authentik_crypto")
    def test_proxy_connectivity(self):
        """Test proxy connectivity over websocket"""
        proxy: ProxyProvider = ProxyProvider.objects.create(
            name=generate_id(),
            authorization_flow=Flow.objects.get(
                slug="default-provider-authorization-implicit-consent"
            ),
            internal_host="http://localhost",
            external_host="http://localhost:9000",
        )
        # Ensure OAuth2 Params are set
        proxy.set_oauth_defaults()
        proxy.save()
        # we need to create an application to actually access the proxy
        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)

        outpost: Outpost = Outpost.objects.create(
            name=generate_id(),
            type=OutpostType.PROXY,
            _config=asdict(OutpostConfig(authentik_host=self.live_server_url, log_level="debug")),
        )
        outpost.providers.add(proxy)
        outpost.build_user_permissions(outpost.user)

        self.run_container(
            image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
            ports={
                "9000": "9000",
            },
            environment={
                "AUTHENTIK_TOKEN": outpost.token.key,
            },
        )

        # Wait until outpost healthcheck succeeds
        healthcheck_retries = 0
        while healthcheck_retries < 50:  # noqa: PLR2004
            if len(outpost.state) > 0:
                state = outpost.state[0]
                if state.last_seen and state.version:
                    break
            healthcheck_retries += 1
            sleep(0.5)

        state = outpost.state
        self.assertGreaterEqual(len(state), 1)

        # Make sure to delete the outpost to remove the container
        outpost.delete()