# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
#
# Copyright (c) 2025-2026 Dmitrii Gagarin aka madgagarin


import ssl
import subprocess
from unittest.mock import AsyncMock, MagicMock, patch

import aiohttp
import pytest
from aiohttp import web

from avtomatika.config import Config
from avtomatika.constants import AUTH_HEADER_WORKER
from avtomatika.engine import OrchestratorEngine
from avtomatika.storage.memory import MemoryStorage


# Helper to run shell commands (reused from test_mtls)
def run_cmd(cmd):
    subprocess.check_call(cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


@pytest.fixture(scope="module")
def pki(tmp_path_factory):
    """Generates a temporary PKI infrastructure for mTLS testing."""
    pki_dir = tmp_path_factory.mktemp("pki_sts")

    # Define filenames
    ca_key = pki_dir / "ca.key"
    ca_crt = pki_dir / "ca.crt"
    server_key = pki_dir / "server.key"
    server_csr = pki_dir / "server.csr"
    server_crt = pki_dir / "server.crt"
    client_key = pki_dir / "client.key"
    client_csr = pki_dir / "client.csr"
    client_crt = pki_dir / "client.crt"

    # Create OpenSSL config file
    openssl_conf = pki_dir / "openssl.cnf"
    with open(openssl_conf, "w") as f:
        f.write("""
[req]
distinguished_name = req_distinguished_name
req_extensions = v3_req

[req_distinguished_name]
commonName = Common Name

[v3_req]
basicConstraints = CA:FALSE
keyUsage = nonRepudiation, digitalSignature, keyEncipherment

[v3_ca]
basicConstraints = critical,CA:TRUE
keyUsage = critical,keyCertSign,cRLSign
subjectKeyIdentifier = hash
authorityKeyIdentifier = keyid:always,issuer
""")

    # 1. Generate CA
    run_cmd(f"openssl genrsa -out {ca_key} 2048")
    run_cmd(
        f"openssl req -x509 -new -nodes -key {ca_key} -sha256 -days 1 "
        f"-out {ca_crt} -subj '/CN=Test CA' -config {openssl_conf} -extensions v3_ca"
    )

    # 2. Generate Server Cert
    run_cmd(f"openssl genrsa -out {server_key} 2048")
    run_cmd(f"openssl req -new -key {server_key} -out {server_csr} -subj '/CN=localhost' -config {openssl_conf}")
    run_cmd(
        f"openssl x509 -req -in {server_csr} -CA {ca_crt} -CAkey {ca_key} "
        f"-CAcreateserial -out {server_crt} -days 1 -sha256 "
        f"-extfile {openssl_conf} -extensions v3_req"
    )

    # 3. Generate Client Cert
    run_cmd(f"openssl genrsa -out {client_key} 2048")
    run_cmd(f"openssl req -new -key {client_key} -out {client_csr} -subj '/CN=worker-sts-test' -config {openssl_conf}")
    run_cmd(
        f"openssl x509 -req -in {client_csr} -CA {ca_crt} -CAkey {ca_key} "
        f"-CAcreateserial -out {client_crt} -days 1 -sha256 "
        f"-extfile {openssl_conf} -extensions v3_req"
    )

    return {
        "ca_crt": str(ca_crt),
        "server_crt": str(server_crt),
        "server_key": str(server_key),
        "client_crt": str(client_crt),
        "client_key": str(client_key),
    }


@pytest.fixture
def mtls_config(pki):
    conf = Config()
    conf.TLS_ENABLED = True
    conf.TLS_CERT_PATH = pki["server_crt"]
    conf.TLS_KEY_PATH = pki["server_key"]
    conf.TLS_CA_PATH = pki["ca_crt"]
    conf.TLS_REQUIRE_CLIENT_CERT = True  # Enforce mTLS for initial connection
    conf.API_PORT = 0
    return conf


@pytest.mark.asyncio
async def test_sts_flow(pki, mtls_config):
    """
    Test the full Security Token Service flow:
    1. Authenticate with mTLS.
    2. Request a temporary access token.
    3. Use the token to authenticate WITHOUT mTLS.
    """
    storage = MemoryStorage()
    engine = OrchestratorEngine(storage, mtls_config)

    engine.setup()
    runner = web.AppRunner(engine.app)
    await runner.setup()

    from rxon.security import create_server_ssl_context

    server_ssl_ctx = create_server_ssl_context(
        cert_path=mtls_config.TLS_CERT_PATH,
        key_path=mtls_config.TLS_KEY_PATH,
        ca_path=mtls_config.TLS_CA_PATH,
        require_client_cert=mtls_config.TLS_REQUIRE_CLIENT_CERT,
    )

    site = web.TCPSite(runner, "localhost", 0, ssl_context=server_ssl_ctx)
    await site.start()

    port = site._server.sockets[0].getsockname()[1] if site._server else 8443

    base_url = f"https://localhost:{port}"
    worker_id = "worker-sts-test"  # Must match cert CN

    # 1. Get Access Token using mTLS
    client_ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=pki["ca_crt"])
    client_ssl_ctx.load_cert_chain(certfile=pki["client_crt"], keyfile=pki["client_key"])

    access_token = None

    async with (
        aiohttp.ClientSession() as session,
        session.post(f"{base_url}/_worker/auth/token", ssl=client_ssl_ctx) as resp,
    ):
        assert resp.status == 200, f"STS request failed: {await resp.text()}"
        data = await resp.json()
        access_token = data.get("access_token")
        assert access_token is not None
        assert data["worker_id"] == worker_id

    # 2. Use Access Token WITHOUT mTLS client cert (but still TLS for encryption)
    # We create a new context that trusts the server CA but has NO client cert.
    # HOWEVER, the server is configured with require_client_cert=True.
    # So we can't connect without a cert at the TCP level!

    # To test token usage, we normally would use it on a non-mTLS port or if client cert is optional.
    # If TLS_REQUIRE_CLIENT_CERT is True, then token usage is redundant for auth (except maybe scope).

    # BUT! STS is useful when:
    # A) Initial bootstrap uses mTLS, then we switch to Token (if cert not required for all requests).
    # B) We use a separate port or ingress for token-based auth.

    # Let's change the test config to verify_client=OPTIONAL, so we can connect without cert
    # but application logic will demand auth.

    await site.stop()  # Stop server to reconfigure

    # Reconfigure server to allow connections without client cert
    server_ssl_ctx.verify_mode = ssl.CERT_OPTIONAL
    site = web.TCPSite(runner, "localhost", port, ssl_context=server_ssl_ctx)
    await site.start()

    # Client context without client cert
    no_cert_ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=pki["ca_crt"])

    async with aiohttp.ClientSession() as session:
        # Try to register using the Token
        headers = {AUTH_HEADER_WORKER: access_token}
        reg_data = {
            "worker_id": worker_id,
            "worker_type": "sts-worker",
            "supported_skills": [{"name": "t1"}],
            "resources": {"max_concurrent_tasks": 1, "cpu_cores": 1, "ram_gb": 1.0},
            "installed_software": {},
            "installed_artifacts": [],
            "capabilities": {"hostname": "h", "ip_address": "127.0.0.1", "cost_per_skill": {}},
        }

        async with session.post(
            f"{base_url}/_worker/workers/register", json=reg_data, headers=headers, ssl=no_cert_ssl_ctx
        ) as resp:
            assert resp.status == 200, f"Token auth failed: {await resp.text()}"
            data = await resp.json()
            assert data["status"] == "registered"

    await site.stop()
    await runner.cleanup()


@pytest.mark.asyncio
async def test_token_hash_cache():
    """Checks that token hashes are cached to save CPU."""
    from avtomatika.security import _TOKEN_HASH_CACHE, verify_worker_auth

    mock_storage = MagicMock()
    mock_storage.verify_worker_access_token = AsyncMock(return_value="w1")
    mock_config = MagicMock()

    _TOKEN_HASH_CACHE.clear()
    token = "token-to-cache"

    # First call
    with patch("avtomatika.security.sha256") as mock_sha:
        mock_sha.return_value.hexdigest.return_value = "hash1"
        await verify_worker_auth(mock_storage, mock_config, token, None, "w1")
        assert mock_sha.called

    # Second call with same token
    with patch("avtomatika.security.sha256") as mock_sha:
        await verify_worker_auth(mock_storage, mock_config, token, None, "w1")
        # Should NOT be called, taken from cache
        assert not mock_sha.called
