"""
Multi-group isolation test: two groups on one QUIC server.

Verifies that:
  - A user in group-a can fetch files from group-a
  - A user in group-a is rejected when requesting group-b
  - A user in both groups can access both
"""

import os
import time
import jwt
import pytest
from pathlib import Path
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives import serialization

from meshbay_common.crypto import generate_gek, pk_to_b64
from meshbay_node.indexer import DirectoryIndexer, GroupIndex
from conftest import one_root
from meshbay_node.transport.quic_server import QuicChunkServer
from meshbay_node.transport.quic_client import QuicChunkClient


@pytest.fixture
def sk_node():
    return Ed25519PrivateKey.generate()

@pytest.fixture
def sk_hub():
    return Ed25519PrivateKey.generate()

@pytest.fixture
def gek_a():
    return generate_gek()

@pytest.fixture
def gek_b():
    return generate_gek()

@pytest.fixture
def dir_a(tmp_path):
    d = tmp_path / "group_a"
    d.mkdir()
    (d / "file_a.txt").write_bytes(b"content from group A " * 100)
    return d

@pytest.fixture
def dir_b(tmp_path):
    d = tmp_path / "group_b"
    d.mkdir()
    (d / "file_b.txt").write_bytes(b"content from group B " * 100)
    return d


def make_jwt(sk_hub, pk_node_b64, groups, user_id="user-001", ttl=3600):
    sk_pem = sk_hub.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.PKCS8,
        serialization.NoEncryption(),
    )
    now = int(time.time())
    return jwt.encode({
        "iss": "test-hub", "sub": user_id,
        "pk_user": pk_node_b64, "hub_id": "test-hub",
        "jti": "test-jti", "iat": now, "exp": now + ttl,
        "groups": groups,
    }, sk_pem, algorithm="EdDSA")


@pytest.fixture
async def multi_group_server(sk_node, sk_hub, gek_a, gek_b, dir_a, dir_b, tmp_path):
    hub_pk_pem = sk_hub.public_key().public_bytes(
        serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo)

    indexer_a = DirectoryIndexer(roots=one_root(dir_a), group_id="group-a", sk_node=sk_node, gek=gek_a)
    await indexer_a.initial_scan()

    indexer_b = DirectoryIndexer(roots=one_root(dir_b), group_id="group-b", sk_node=sk_node, gek=gek_b)
    await indexer_b.initial_scan()

    groups = {
        "group-a": {"gek": gek_a, "roots": dir_a, "index": indexer_a.index},
        "group-b": {"gek": gek_b, "roots": dir_b, "index": indexer_b.index},
    }

    cert_path = tmp_path / "node.crt"
    key_path  = tmp_path / "node.key"

    server = QuicChunkServer(
        sk_node=sk_node, hub_pk_pem=hub_pk_pem,
        gek=gek_a, roots=one_root(dir_a), index=indexer_a.index,
        host="127.0.0.1", port=19200,
        cert_path=cert_path, key_path=key_path,
        groups=groups,
    )
    await server.start()
    yield server, indexer_a, indexer_b
    await server.stop()


@pytest.mark.asyncio
async def test_user_can_access_own_group(
    multi_group_server, sk_node, sk_hub, gek_a,
):
    """User in group-a can fetch index and chunks from group-a."""
    server, indexer_a, _ = multi_group_server
    token = make_jwt(sk_hub, pk_to_b64(sk_node.public_key()), groups=["group-a"])

    async with QuicChunkClient(
        host="127.0.0.1", port=19200,
        jwt_token=token, gek=gek_a,
        pk_node_b64=pk_to_b64(sk_node.public_key()),
        group_id="group-a",
    ) as client:
        wire = await client.fetch_index()
        recovered = GroupIndex.deserialize(wire, sk_node=sk_node, gek=gek_a)
        assert recovered.count == 1

        entry = recovered.entries[0]
        assert entry.name == "file_a.txt"
        chunk = await client.fetch_chunk(entry.id, chunk_index=0)
        assert chunk == b"content from group A " * 100


@pytest.mark.asyncio
async def test_user_rejected_from_other_group(
    multi_group_server, sk_node, sk_hub, gek_b,
):
    """User in group-a only is rejected when requesting group-b."""
    server, _, _ = multi_group_server
    token = make_jwt(sk_hub, pk_to_b64(sk_node.public_key()), groups=["group-a"])

    with pytest.raises(ConnectionError, match="rejected"):
        async with QuicChunkClient(
            host="127.0.0.1", port=19200,
            jwt_token=token, gek=gek_b,
            pk_node_b64=pk_to_b64(sk_node.public_key()),
            group_id="group-b",
        ) as client:
            await client.fetch_index()


@pytest.mark.asyncio
async def test_dual_group_user_accesses_both(
    multi_group_server, sk_node, sk_hub, gek_a, gek_b,
):
    """User in both groups can access either group's files."""
    server, indexer_a, indexer_b = multi_group_server
    token = make_jwt(sk_hub, pk_to_b64(sk_node.public_key()), groups=["group-a", "group-b"])

    async with QuicChunkClient(
        host="127.0.0.1", port=19200,
        jwt_token=token, gek=gek_a,
        pk_node_b64=pk_to_b64(sk_node.public_key()),
        group_id="group-a",
    ) as client_a:
        wire_a = await client_a.fetch_index()
        idx_a = GroupIndex.deserialize(wire_a, sk_node=sk_node, gek=gek_a)
        assert idx_a.entries[0].name == "file_a.txt"

    async with QuicChunkClient(
        host="127.0.0.1", port=19200,
        jwt_token=token, gek=gek_b,
        pk_node_b64=pk_to_b64(sk_node.public_key()),
        group_id="group-b",
    ) as client_b:
        wire_b = await client_b.fetch_index()
        idx_b = GroupIndex.deserialize(wire_b, sk_node=sk_node, gek=gek_b)
        assert idx_b.entries[0].name == "file_b.txt"