"""Authentication helpers combining crypto with autogenerated API functions."""

from __future__ import annotations

import base64
import time

from arbi_client.api.user import login_user, register_user
from arbi_client.client import Client
from arbi_client.models.login_request import LoginRequest
from arbi_client.models.register_request import RegisterRequest
from arbi_client.sdk._types import SessionState, UserIdentity, WorkspaceInfo
from arbi_client.sdk.crypto import create_login_signature, generate_keypair

# ---------------------------------------------------------------------------
# Agent identity helpers
# ---------------------------------------------------------------------------


def _build_agent_identity(
    agent_email: str,
    signing_key: bytes,
) -> UserIdentity:
    """Build a UserIdentity for a persistent agent from its Ed25519 signing key.

    Derives X25519 encryption keys from the Ed25519 key (same curve conversion
    used during registration).
    """
    from nacl.signing import SigningKey

    sk = SigningKey(signing_key[:32])  # seed is first 32 bytes
    signing_pub_b64 = base64.b64encode(bytes(sk.verify_key)).decode()
    x25519_priv = sk.to_curve25519_private_key()
    enc_pub_b64 = base64.b64encode(bytes(x25519_priv.public_key)).decode()
    enc_priv_bytes = bytes(x25519_priv)

    return UserIdentity(
        email=agent_email,
        external_id="",  # filled in from login response
        given_name="",
        family_name="",
        signing_private_key=bytes(sk),
        signing_public_key_b64=signing_pub_b64,
        encryption_private_key=enc_priv_bytes,
        encryption_public_key_b64=enc_pub_b64,
    )


# ---------------------------------------------------------------------------
# Register
# ---------------------------------------------------------------------------


def register(
    client: Client,
    *,
    email: str,
    password: str,
    given_name: str,
    family_name: str,
    deployment_domain: str,
    verification_credential: str,
) -> UserIdentity:
    """Register a new user and return their identity (keys derived from password).

    Does **not** log in — call :func:`login` or :func:`login_with_identity` afterwards.
    """
    signing_pub_b64, signing_priv, enc_pub_b64, enc_priv = generate_keypair(password, email, deployment_domain)

    register_user.sync(
        client=client,
        body=RegisterRequest(
            email=email,
            signing_key=signing_pub_b64,
            verification_credential=verification_credential,
            given_name=given_name,
            family_name=family_name,
        ),
    )

    # Login to get the user's external_id
    sig_b64, timestamp = create_login_signature(email, signing_priv)
    login_resp = login_user.sync(
        client=client,
        body=LoginRequest(email=email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login after registration failed: {login_resp}")

    return UserIdentity(
        email=email,
        external_id=login_resp.user.external_id,
        given_name=given_name,
        family_name=family_name,
        signing_private_key=signing_priv,
        signing_public_key_b64=signing_pub_b64,
        encryption_private_key=enc_priv,
        encryption_public_key_b64=enc_pub_b64,
    )


async def aregister(
    client: Client,
    *,
    email: str,
    password: str,
    given_name: str,
    family_name: str,
    deployment_domain: str,
    verification_credential: str,
) -> UserIdentity:
    """Async variant of :func:`register`."""
    signing_pub_b64, signing_priv, enc_pub_b64, enc_priv = generate_keypair(password, email, deployment_domain)

    await register_user.asyncio(
        client=client,
        body=RegisterRequest(
            email=email,
            signing_key=signing_pub_b64,
            verification_credential=verification_credential,
            given_name=given_name,
            family_name=family_name,
        ),
    )

    sig_b64, timestamp = create_login_signature(email, signing_priv)
    login_resp = await login_user.asyncio(
        client=client,
        body=LoginRequest(email=email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login after registration failed: {login_resp}")

    return UserIdentity(
        email=email,
        external_id=login_resp.user.external_id,
        given_name=given_name,
        family_name=family_name,
        signing_private_key=signing_priv,
        signing_public_key_b64=signing_pub_b64,
        encryption_private_key=enc_priv,
        encryption_public_key_b64=enc_pub_b64,
    )


# ---------------------------------------------------------------------------
# Login
# ---------------------------------------------------------------------------


def _extract_workspaces(login_resp) -> list[WorkspaceInfo]:
    """Extract workspace info from login response for header caching."""
    from arbi_client.types import UNSET

    raw = getattr(login_resp, "workspaces", UNSET)
    if raw is UNSET or raw is None:
        return []
    return [
        WorkspaceInfo(
            external_id=ws.additional_properties.get("external_id", ""),
            wrapped_key=ws.additional_properties.get("wrapped_key"),
        )
        for ws in raw
        if hasattr(ws, "additional_properties")
    ]


def login(
    client: Client,
    *,
    email: str,
    password: str,
    deployment_domain: str,
) -> tuple[UserIdentity, SessionState, list[WorkspaceInfo]]:
    """Derive keys from password, log in, and return identity + session state + workspaces."""
    signing_pub_b64, signing_priv, enc_pub_b64, enc_priv = generate_keypair(password, email, deployment_domain)

    sig_b64, timestamp = create_login_signature(email, signing_priv)

    login_resp = login_user.sync(
        client=client,
        body=LoginRequest(email=email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login failed: {login_resp}")

    identity = UserIdentity(
        email=email,
        external_id=login_resp.user.external_id,
        given_name=login_resp.user.given_name or "",
        family_name=login_resp.user.family_name or "",
        signing_private_key=signing_priv,
        signing_public_key_b64=signing_pub_b64,
        encryption_private_key=enc_priv,
        encryption_public_key_b64=enc_pub_b64,
    )

    session = SessionState(
        access_token=login_resp.access_token,
        session_public_key_b64=login_resp.session_key,
    )

    return identity, session, _extract_workspaces(login_resp)


async def alogin(
    client: Client,
    *,
    email: str,
    password: str,
    deployment_domain: str,
) -> tuple[UserIdentity, SessionState, list[WorkspaceInfo]]:
    """Async variant of :func:`login`."""
    signing_pub_b64, signing_priv, enc_pub_b64, enc_priv = generate_keypair(password, email, deployment_domain)

    sig_b64, timestamp = create_login_signature(email, signing_priv)

    login_resp = await login_user.asyncio(
        client=client,
        body=LoginRequest(email=email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login failed: {login_resp}")

    identity = UserIdentity(
        email=email,
        external_id=login_resp.user.external_id,
        given_name=login_resp.user.given_name or "",
        family_name=login_resp.user.family_name or "",
        signing_private_key=signing_priv,
        signing_public_key_b64=signing_pub_b64,
        encryption_private_key=enc_priv,
        encryption_public_key_b64=enc_pub_b64,
    )

    session = SessionState(
        access_token=login_resp.access_token,
        session_public_key_b64=login_resp.session_key,
    )

    return identity, session, _extract_workspaces(login_resp)


# ---------------------------------------------------------------------------
# Login with existing identity (re-auth without password)
# ---------------------------------------------------------------------------


def login_with_identity(
    client: Client,
    *,
    identity: UserIdentity,
) -> tuple[SessionState, list[WorkspaceInfo]]:
    """Log in using an existing :class:`UserIdentity` (skips key derivation)."""
    sig_b64, timestamp = create_login_signature(identity.email, identity.signing_private_key)

    login_resp = login_user.sync(
        client=client,
        body=LoginRequest(email=identity.email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login failed: {login_resp}")

    session = SessionState(
        access_token=login_resp.access_token,
        session_public_key_b64=login_resp.session_key,
    )
    return session, _extract_workspaces(login_resp)


async def alogin_with_identity(
    client: Client,
    *,
    identity: UserIdentity,
) -> tuple[SessionState, list[WorkspaceInfo]]:
    """Async variant of :func:`login_with_identity`."""
    sig_b64, timestamp = create_login_signature(identity.email, identity.signing_private_key)

    login_resp = await login_user.asyncio(
        client=client,
        body=LoginRequest(email=identity.email, signature=sig_b64, timestamp=timestamp),
    )
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Login failed: {login_resp}")

    session = SessionState(
        access_token=login_resp.access_token,
        session_public_key_b64=login_resp.session_key,
    )
    return session, _extract_workspaces(login_resp)


# ---------------------------------------------------------------------------
# Agent login (persistent agent re-authentication)
# ---------------------------------------------------------------------------


def _prepare_agent_login(agent_email: str, signing_key: bytes) -> tuple[LoginRequest, UserIdentity]:
    """Build the login request body and identity for an agent login."""
    from nacl.signing import SigningKey

    sk = SigningKey(signing_key[:32])
    timestamp = int(time.time())
    message = f"{agent_email}|{timestamp}".encode()
    sig_b64 = base64.b64encode(sk.sign(message).signature).decode()

    body = LoginRequest(email=agent_email, signature=sig_b64, timestamp=timestamp)
    identity = _build_agent_identity(agent_email, signing_key)
    return body, identity


def _finish_agent_login(login_resp, identity: UserIdentity) -> tuple[UserIdentity, SessionState, list[WorkspaceInfo]]:
    """Validate response and build return tuple for agent login."""
    if login_resp is None or not hasattr(login_resp, "access_token"):
        raise RuntimeError(f"Agent login failed: {login_resp}")
    session = SessionState(
        access_token=login_resp.access_token,
        session_public_key_b64=login_resp.session_key,
    )
    return identity, session, _extract_workspaces(login_resp)


def agent_login(
    client: Client,
    *,
    agent_email: str,
    signing_key: bytes,
) -> tuple[UserIdentity, SessionState, list[WorkspaceInfo]]:
    """Log in as a persistent agent using Ed25519 signature.

    The agent's X25519 encryption key is derived from its Ed25519 signing key,
    enabling the same workspace key unwrap flow as regular users.

    Args:
        client: Unauthenticated API client.
        agent_email: Agent's email address (e.g. ``agentname.usr-XXXXXXXX@deploymentdomain``).
        signing_key: 64-byte Ed25519 private key (or 32-byte seed).

    Returns:
        ``(identity, session, workspaces)`` — same shape as :func:`login`.
    """
    body, identity = _prepare_agent_login(agent_email, signing_key)
    login_resp = login_user.sync(client=client, body=body)
    return _finish_agent_login(login_resp, identity)


async def aagent_login(
    client: Client,
    *,
    agent_email: str,
    signing_key: bytes,
) -> tuple[UserIdentity, SessionState, list[WorkspaceInfo]]:
    """Async variant of :func:`agent_login`."""
    body, identity = _prepare_agent_login(agent_email, signing_key)
    login_resp = await login_user.asyncio(client=client, body=body)
    return _finish_agent_login(login_resp, identity)
