import logging
import ssl
import sys
from typing import Any, Dict, Optional

from keylime import api_version as keylime_api_version
from keylime import keylime_logging
from keylime.ip_util import bracketize_ipv6
from keylime.requests_client import RequestsClient

if sys.version_info >= (3, 8):
    from typing import TypedDict
else:
    from typing_extensions import TypedDict

if sys.version_info >= (3, 11):
    from typing import NotRequired
else:
    from typing_extensions import NotRequired


class RegistrarData(TypedDict):
    ip: Optional[str]
    port: Optional[str]
    regcount: int
    mtls_cert: Optional[str]
    aik_tpm: str
    ek_tpm: str
    ekcert: Optional[str]
    provider_keys: NotRequired[Dict[str, str]]


logger = keylime_logging.init_logging("registrar_client")
API_VERSION = keylime_api_version.current_version()

MANDATORY_FIELDS = ["aik_tpm", "regcount", "ek_tpm", "ip", "port"]


def check_mandatory_fields(results: Dict[str, Any]) -> bool:
    for field in MANDATORY_FIELDS:
        if field not in results:
            logger.critical("Error: did not receive %s from Registrar Server.", field)
            return False
    return True


def getVersions(
    registrar_ip: str,
    registrar_port: str,
    tls_context: Optional[ssl.SSLContext],
) -> Optional[Dict[str, Any]]:
    """
    Fetch supported API versions from the registrar's /version endpoint.

    :returns: JSON structure containing version info, or None on failure
    """
    try:
        client = RequestsClient(f"{bracketize_ipv6(registrar_ip)}:{registrar_port}", True, tls_context=tls_context)
        response = client.get("/version")

        if response.status_code != 200:
            logger.warning("Failed to get versions from registrar: %s", response.status_code)
            return None

        response_body: Dict[str, Any] = response.json()
        if "results" not in response_body:
            logger.warning("Unexpected response format from registrar /version endpoint")
            return None

        results: Dict[str, Any] = response_body["results"]
        return results

    except Exception as e:
        logger.warning("Error fetching versions from registrar: %s", e)
        return None


def getData(
    registrar_ip: str,
    registrar_port: str,
    agent_id: str,
    tls_context: Optional[ssl.SSLContext],
    api_version: Optional[str] = None,
) -> Optional[RegistrarData]:
    """
    Get the agent data from the registrar.

    This is called by the tenant code

    :param api_version: Optional API version to use. If None, uses current version.
    :returns: JSON structure containing the agent data
    """
    # make absolutely sure you don't ask for data that contains AIK keys unauthenticated
    if not tls_context:
        raise Exception("It is unsafe to use this interface to query AIKs without server-authenticated TLS.")

    version = api_version or API_VERSION
    response = None
    try:
        client = RequestsClient(f"{bracketize_ipv6(registrar_ip)}:{registrar_port}", True, tls_context=tls_context)
        response = client.get(f"/v{version}/agents/{agent_id}")
        response_body = response.json()

        if response.status_code == 404:
            logger.critical(
                "Error: could not get agent %s data from Registrar Server: %s", agent_id, response.status_code
            )
            keylime_logging.log_http_response(logger, logging.CRITICAL, response_body)
            return None

        if response.status_code != 200:
            logger.critical("Error: unexpected http response code from Registrar Server: %s", response.status_code)
            keylime_logging.log_http_response(logger, logging.CRITICAL, response_body)
            return None

        # Check for all values that are consumed by other parts of Keylime
        if "results" not in response_body:
            logger.critical("Error: unexpected http response body from Registrar Server: %s", response.status_code)
            return None

        if not check_mandatory_fields(response_body["results"]):
            return None

        r = response_body["results"]
        res: RegistrarData = {
            "aik_tpm": r["aik_tpm"],
            "regcount": r["regcount"],
            "ek_tpm": r["ek_tpm"],
            "ip": r["ip"],
            "port": r["port"],
            "mtls_cert": r.get("mtls_cert"),
            "ekcert": r.get("ekcert"),
        }
        if "provider_keys" in r:
            res["provider_keys"] = r["provider_keys"]

        return res

    except AttributeError as e:
        if response and response.status_code == 503:
            logger.critical("Error: the registrar is not available at %s:%s", registrar_ip, registrar_port)
        else:
            logger.exception(e)

    except Exception as e:
        logger.exception(e)

    return None


def doRegistrarDelete(
    registrar_ip: str,
    registrar_port: str,
    agent_id: str,
    tls_context: Optional[ssl.SSLContext],
    api_version: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Delete the given agent from the registrar.

    This is called by the tenant code

    :param api_version: Optional API version to use. If None, uses current version.
    :returns: The request response body
    """
    version = api_version or API_VERSION
    client = RequestsClient(f"{bracketize_ipv6(registrar_ip)}:{registrar_port}", True, tls_context=tls_context)
    response = client.delete(f"/v{version}/agents/{agent_id}")
    response_body: Dict[str, Any] = response.json()

    if response.status_code == 200:
        logger.debug("Registrar deleted.")
    else:
        logger.warning("Status command response: %s Unexpected response from registrar.", response.status_code)
        keylime_logging.log_http_response(logger, logging.WARNING, response_body)

    return response_body


def doRegistrarList(
    registrar_ip: str,
    registrar_port: str,
    tls_context: Optional[ssl.SSLContext],
    api_version: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """
    Get the list of registered agents from the registrar.

    This is called by the tenant code

    :param api_version: Optional API version to use. If None, uses current version.
    :returns: The request response body
    """
    version = api_version or API_VERSION
    client = RequestsClient(f"{bracketize_ipv6(registrar_ip)}:{registrar_port}", True, tls_context=tls_context)
    response = client.get(f"/v{version}/agents/")
    response_body: Dict[str, Any] = response.json()

    if response.status_code != 200:
        logger.warning("Registrar returned: %s Unexpected response from registrar.", response.status_code)
        keylime_logging.log_http_response(logger, logging.WARNING, response_body)
        return None

    return response_body
