#!/usr/bin/env python3
# MIT License
#
# Author: Komal Thareja (kthare10@renci.org)

from __future__ import annotations

import json
import logging
import time

from typing import Any, Dict, List, Literal, Optional, Union

from fim.user.topology import AdvertizedTopology, Topology

from fabrictestbed.external_api.credmgr_client import CredmgrClient
from fabrictestbed.external_api.orchestrator_client import OrchestratorClient, SliverDTO, SliceDTO, PoaDataDTO
from fabrictestbed.external_api.core_api import CoreApi
from fabrictestbed.topology_query_api import TopologyQueryAPI

# =========================
# FabricManagerV2 façade
# =========================

class FabricManagerV2(TopologyQueryAPI):
    """
    One-stop façade over CredmgrClient + OrchestratorClient.

    - All methods accept `return_fmt="dict"|"dto"`.
    - For orchestrator calls, pass `id_token`.
    - Token helpers available via `credmgr`.
    """

    def __init__(
        self,
        *,
        credmgr_host: str,
        orchestrator_host: str,
        core_api_host: Optional[str] = None,
        logger: Optional[logging.Logger] = None,
        http_debug: bool = False,
    ):
        super().__init__(orchestrator_host=orchestrator_host, logger=logger, http_debug=http_debug)
        self.credmgr = CredmgrClient(credmgr_host, http_debug=http_debug, logger=self.log)
        self.orch = OrchestratorClient(orchestrator_host, http_debug=http_debug, logger=self.log)
        self.core_api_host = core_api_host

    # -------- Token helpers (proxy) --------

    def tokens_create(self, **kwargs) -> List[Dict[str, Any]]:
        """
        Create new FABRIC tokens.

        Proxies to CredmgrClient.create with dict return format for MCP compatibility.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.create.

        Returns:
            List of token dictionaries with token details.
        """
        # Same signature as CredmgrV3.create; return dicts by default (MCP friendly)
        return self.credmgr.create(return_fmt="dict", **kwargs)  # type: ignore[arg-type]

    def tokens_refresh(self, **kwargs) -> List[Dict[str, Any]]:
        """
        Refresh existing FABRIC tokens.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.refresh.

        Returns:
            List of refreshed token dictionaries.
        """
        return self.credmgr.refresh(return_fmt="dict", **kwargs)  # type: ignore[arg-type]

    def tokens_validate(self, *, id_token: str) -> Dict[str, Any]:
        """
        Validate a FABRIC ID token.

        Args:
            id_token: The FABRIC ID token to validate.

        Returns:
            Dictionary containing validation result and token details.
        """
        return self.credmgr.validate(id_token=id_token, return_fmt="dict")  # type: ignore[return-value]

    def tokens_list(self, **kwargs) -> List[Dict[str, Any]]:
        """
        List all FABRIC tokens for the user.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.tokens.

        Returns:
            List of token dictionaries.
        """
        return self.credmgr.tokens(return_fmt="dict", **kwargs)  # type: ignore[arg-type]

    def tokens_revoke(self, **kwargs) -> None:
        """
        Revoke a FABRIC token.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.revoke.
        """
        return self.credmgr.revoke(**kwargs)

    def tokens_delete(self, **kwargs) -> None:
        """
        Delete a FABRIC token.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.delete.
        """
        return self.credmgr.delete(**kwargs)

    def tokens_delete_all(self, **kwargs) -> None:
        """
        Delete all FABRIC tokens for the user.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.delete_all.
        """
        return self.credmgr.delete_all(**kwargs)

    def tokens_revoke_list(self, **kwargs) -> List[str]:
        """
        Get list of revoked tokens.

        Args:
            **kwargs: Keyword arguments passed to CredmgrClient.revoke_list.

        Returns:
            List of revoked token identifiers.
        """
        return self.credmgr.revoke_list(**kwargs)

    # -------- Orchestrator helpers --------
    def create_slice(
        self,
        *,
        id_token: str,
        name: str,
        graph_model: str,
        ssh_keys: List[str],
        lifetime: Optional[int] = None,
        lease_start_time: Optional[str] = None,
        lease_end_time: Optional[str] = None,
        return_fmt: Literal["dict", "dto"] = "dict",
    ) -> List[Union[Dict[str, Any], SliverDTO]]:
        """
        Create a new FABRIC slice.

        Args:
            id_token: FABRIC ID token for authentication.
            name: Name of the slice to create.
            graph_model: Slice topology graph model (GRAPHML, JSON, etc.).
            ssh_keys: List of SSH public keys for slice access.
            lifetime: Optional slice lifetime in days.
            lease_start_time: Optional lease start time (UTC format).
            lease_end_time: Optional lease end time (UTC format).
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of sliver dictionaries or DTO objects representing the created slice resources.
        """
        return self.orch.create(
            token=id_token,
            slice_name=name,
            slice_graph=graph_model,
            ssh_keys=ssh_keys,
            lifetime=lifetime,
            lease_start_time=lease_start_time,
            lease_end_time=lease_end_time,
            return_fmt=return_fmt,
        )

    def modify_slice(
        self, *, id_token: str, slice_id: str, graph_model: str, return_fmt: Literal["dict", "dto"] = "dict"
    ) -> List[Union[Dict[str, Any], SliverDTO]]:
        """
        Modify an existing FABRIC slice topology.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice to modify.
            graph_model: Updated slice topology graph model.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of sliver dictionaries or DTO objects with modification results.
        """
        return self.orch.modify(token=id_token, slice_id=slice_id, slice_graph=graph_model, return_fmt=return_fmt)

    def accept_modify(
        self, *, id_token: str, slice_id: str, return_fmt: Literal["dict", "dto"] = "dict"
    ) -> Union[Dict[str, Any], SliceDTO]:
        """
        Accept pending slice modifications.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice with pending modifications.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            Slice dictionary or DTO object with updated state.
        """
        return self.orch.modify_accept(token=id_token, slice_id=slice_id, return_fmt=return_fmt)

    def renew_slice(self, *, id_token: str, slice_id: str, lease_end_time: str) -> None:
        """
        Renew a FABRIC slice lease.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice to renew.
            lease_end_time: New lease end time (UTC format).
        """
        return self.orch.renew(token=id_token, slice_id=slice_id, new_lease_end_time=lease_end_time)

    def delete_slice(self, *, id_token: str, slice_id: Optional[str] = None) -> None:
        """
        Delete a FABRIC slice.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: Optional UUID of the slice to delete.
        """
        return self.orch.delete(token=id_token, slice_id=slice_id)

    def list_slices(
        self,
        *,
        id_token: str,
        states: Optional[List[str]] = None,
        exclude_states: Optional[List[str]] = None,
        name: Optional[str] = None,
        search: Optional[str] = None,
        exact_match: bool = False,
        as_self: bool = True,
        limit: int = 200,
        offset: int = 0,
        return_fmt: Literal["dict", "dto"] = "dict",
    ) -> List[Union[Dict[str, Any], SliceDTO]]:
        """
        List FABRIC slices with optional filtering.

        Args:
            id_token: FABRIC ID token for authentication.
            states: Optional list of slice states to include (e.g., ["StableError", "StableOK"]). Allowed values: (Nascent, Configuring, StableOK,
                     StableError, ModifyOK, ModifyError, Closing, Dead).
            exclude_states: Optional list of slice states to exclude (e.g., for fetching active slices set exclude_states=["Closing", "Dead"]).
            name: Optional slice name filter.
            search: Optional search string for slice names.
            exact_match: If True, match slice name exactly; if False, use substring match.
            as_self: If True, list only user's own slices; if False, list all accessible slices.
            limit: Maximum number of slices to return (default: 200).
            offset: Pagination offset (default: 0).
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of slice dictionaries or DTO objects.
        """
        return self.orch.slices(
            token=id_token,
            includes=states,
            excludes=exclude_states,
            name=name,
            search=search,
            exact_match=exact_match,
            as_self=as_self,
            limit=limit,
            offset=offset,
            return_fmt=return_fmt,
        )

    def get_slice(
        self,
        *,
        id_token: str,
        slice_id: str,
        graph_format: Literal["GRAPHML", "JSON_NODELINK", "CYTOSCAPE", "NONE"] = "GRAPHML",
        as_self: bool = True,
        return_fmt: Literal["dict", "dto"] = "dict",
    ) -> Union[Dict[str, Any], SliceDTO]:
        """
        Get details of a specific FABRIC slice.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice to retrieve.
            graph_format: Format for the slice topology graph (GRAPHML, JSON_NODELINK, CYTOSCAPE, or NONE).
            as_self: If True, retrieve as owner; if False, retrieve with delegated access.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            Slice dictionary or DTO object with full details.
        """
        return self.orch.get_slice(token=id_token, slice_id=slice_id, graph_format=graph_format,
                                   as_self=as_self, return_fmt=return_fmt)

    def list_slivers(
        self, *, id_token: str, slice_id: str, as_self: bool = True, return_fmt: Literal["dict", "dto"] = "dict"
    ) -> List[Union[Dict[str, Any], SliverDTO]]:
        """
        List all slivers (resource allocations) in a slice.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice containing the slivers.
            as_self: If True, list as owner; if False, list with delegated access.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of sliver dictionaries or DTO objects.
        """
        return self.orch.slivers(token=id_token, slice_id=slice_id, as_self=as_self, return_fmt=return_fmt)

    def get_sliver(
        self, *, id_token: str, slice_id: str, sliver_id: str, as_self: bool = True, return_fmt: Literal["dict", "dto"] = "dict"
    ) -> List[Union[Dict[str, Any], SliverDTO]]:
        """
        Get details of a specific sliver.

        Args:
            id_token: FABRIC ID token for authentication.
            slice_id: UUID of the slice containing the sliver.
            sliver_id: UUID of the sliver to retrieve.
            as_self: If True, retrieve as owner; if False, retrieve with delegated access.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List containing the sliver dictionary or DTO object.
        """
        return self.orch.slivers(token=id_token, slice_id=slice_id, sliver_id=sliver_id, as_self=as_self,
                                 return_fmt=return_fmt)

    def poa_create(
        self,
        *,
        id_token: str,
        sliver_id: str,
        operation: Literal["cpuinfo", "numainfo", "cpupin", "numatune", "reboot", "addkey", "removekey", "rescan"],
        vcpu_cpu_map: Optional[List[Dict[str, str]]] = None,
        node_set: Optional[List[str]] = None,
        keys: Optional[List[Dict[str, str]]] = None,
        bdf: Optional[List[str]] = None,
        return_fmt: Literal["dict", "dto"] = "dict",
    ) -> List[Union[Dict[str, Any], PoaDataDTO]]:
        """
        Create a Perform Operational Action (POA) on a sliver.

        POAs allow runtime operations on deployed resources like CPU pinning,
        NUMA tuning, reboots, SSH key management, and device rescanning.

        Args:
            id_token: FABRIC ID token for authentication.
            sliver_id: UUID of the sliver to perform the action on.
            operation: The POA operation to perform (cpuinfo, numainfo, cpupin, numatune,
                      reboot, addkey, removekey, rescan).
            vcpu_cpu_map: Optional list of vCPU-to-CPU mapping dictionaries for cpupin operation.
            node_set: Optional list of NUMA node identifiers for numatune operation.
            keys: Optional list of SSH key dictionaries for addkey/removekey operations.
            bdf: Optional list of Bus:Device.Function identifiers for rescan operation.
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of POA result dictionaries or DTO objects.
        """
        return self.orch.poa(
            token=id_token,
            sliver_id=sliver_id,
            operation=operation,
            vcpu_cpu_map=vcpu_cpu_map,
            node_set=node_set,
            keys=keys,
            bdf=bdf,
            return_fmt=return_fmt,
        )

    def poa_get(
        self,
        *,
        id_token: str,
        sliver_id: Optional[str] = None,
        poa_id: Optional[str] = None,
        limit: int = 200,
        offset: int = 0,
        return_fmt: Literal["dict", "dto"] = "dict",
    ) -> List[Union[Dict[str, Any], PoaDataDTO]]:
        """
        Retrieve POA (Perform Operational Action) status and results.

        Args:
            id_token: FABRIC ID token for authentication.
            sliver_id: Optional UUID of the sliver to filter POAs by.
            poa_id: Optional UUID of a specific POA to retrieve.
            limit: Maximum number of POAs to return (default: 200).
            offset: Pagination offset (default: 0).
            return_fmt: Return format - "dict" for dictionaries or "dto" for DTO objects.

        Returns:
            List of POA dictionaries or DTO objects with status and results.
        """
        return self.orch.get_poas(
            token=id_token, sliver_id=sliver_id, poa_id=poa_id, limit=limit, offset=offset, return_fmt=return_fmt
        )

    # -------- Storage helpers (Core API) --------
    def list_storage(
        self, *, id_token: str, offset: int = 0, limit: int = 200
    ) -> List[Dict[str, Any]]:
        """
        List all storage volumes.

        :param id_token: FABRIC ID token for authentication.
        :param offset: Pagination offset (default: 0).
        :param limit: Maximum number of records to fetch (default: 200).
        :return: List of storage volume dictionaries.
        """
        if not self.core_api_host:
            raise ValueError("core_api_host must be provided during initialization to use storage methods")
        core_api = CoreApi(core_api_host=self.core_api_host, token=id_token)
        return core_api.list_storage(offset=offset, limit=limit)

    def get_storage(self, *, id_token: str, uuid: str) -> List[Dict[str, Any]]:
        """
        Get a specific storage volume by UUID.

        :param id_token: FABRIC ID token for authentication.
        :param uuid: Storage volume UUID.
        :return: Storage volume details.
        """
        if not self.core_api_host:
            raise ValueError("core_api_host must be provided during initialization to use storage methods")
        core_api = CoreApi(core_api_host=self.core_api_host, token=id_token)
        return core_api.get_storage(uuid=uuid)



if __name__ == "__main__":
    id_token = ""
    mgr = FabricManagerV2(credmgr_host="cm.fabric-testbed.net",
                          orchestrator_host="orchestrator.fabric-testbed.net",)

    start = time.perf_counter()
    sites = mgr.query_sites(id_token=id_token)
    print(json.dumps(sites, indent=2))
    print("Query sites time:", time.perf_counter() - start)


    start = time.perf_counter()
    hosts = mgr.query_hosts(id_token=id_token)
    print(json.dumps(hosts, indent=2))
    print("Query hosts time:", time.perf_counter() - start)

    start = time.perf_counter()
    fps = mgr.query_facility_ports(id_token=id_token)
    print(json.dumps(fps, indent=2))
    print("Query Facility Ports time:", time.perf_counter() - start)

    start = time.perf_counter()
    fps = mgr.query_links(id_token=id_token)
    print(json.dumps(fps, indent=2))
    print("Query Links time:", time.perf_counter() - start)
