#!/usr/bin/env python
# -*- coding: UTF-8 -*-
#
# Copyright 2021-2023 NXP
#
# SPDX-License-Identifier: BSD-3-Clause

"""Module for general utilities used by TP applications."""
import itertools
import os
from typing import Any, Callable, Dict, List, Optional, Union

import click
import colorama
import prettytable
from typing_extensions import Literal

from spsdk import SPSDKError
from spsdk.apps.utils import FC
from spsdk.tp import TP_DATA_FOLDER, TP_SCH_FILE, SPSDKTpError
from spsdk.tp.tp_intf import TpDevInterface
from spsdk.tp.tphost import TrustProvisioningHost
from spsdk.tp.utils import (
    TpIntfDescription,
    get_tp_device_class,
    get_tp_device_types,
    get_tp_target_class,
    get_tp_target_types,
    scan_tp_devices,
    scan_tp_targets,
    single_tp_device_adapter,
    single_tp_target_adapter,
)
from spsdk.utils.misc import find_file, load_binary, load_configuration
from spsdk.utils.schema_validator import ValidationSchemas, check_config


class TPBaseConfig:
    """Base class for TP app configs."""

    SCHEMA_FILE_DIR = TP_DATA_FOLDER
    SCHEMA_MEMBERS: List[str] = []

    def __init__(self, config_data: dict, config_dir: Optional[str] = None) -> None:
        """Initialize the basic configuration.

        :param config_data: Initial configuration data
        :param config_dir: Path to configuration file
        """
        self.config_data = config_data
        self.config_dir = config_dir

    def _validate(self) -> None:
        """Validate configuration data using appropriate validation schema."""
        schema_cfg = ValidationSchemas.get_schema_file(TP_SCH_FILE)

        # Get this app type scheme pieces
        sch_list = [schema_cfg[x] for x in self.SCHEMA_MEMBERS]
        # First check common settings
        check_config(
            config=self.config_data,
            schemas=sch_list,
            search_paths=[self.config_dir] if self.config_dir else None,
        )

        # Add device schemas
        if "device" in self.SCHEMA_MEMBERS:
            sch_list.extend(
                get_tp_device_class(self.config_data["tp_device"]).get_validation_schemas()
            )
        # Optionally add also target schemas
        if "target" in self.SCHEMA_MEMBERS:
            sch_list.extend(
                get_tp_target_class(self.config_data["tp_target"]).get_validation_schemas()
            )

        check_config(
            config=self.config_data,
            schemas=sch_list,
            search_paths=[self.config_dir] if self.config_dir else None,
        )

    @staticmethod
    def _get_config_data(config_file_path: Optional[str] = None) -> Dict[str, Any]:
        """Setup initial configuration data."""
        if config_file_path:
            return load_configuration(config_file_path)
        return {"timeout": 60}

    @property
    def tp_device(self) -> str:
        """Trust Provisioning device."""
        return self.config_data["tp_device"]

    @property
    def tp_device_parameter(self) -> Dict[str, Any]:
        """Trust Provisioning device parameters."""
        return self.config_data["tp_device_parameter"]

    @property
    def timeout(self) -> int:
        """Timeout."""
        return self.config_data.get("timeout", 60)

    @property
    def family(self) -> str:
        """Target chip family."""
        return self.config_data["family"]


class TPHostConfig(TPBaseConfig):
    """Configuration class for TPHost app."""

    SCHEMA_MEMBERS = [
        "family",
        "provisioning_firmware",
        "oem_firmware",
        "tp_timeout",
        "audit_log",
        "device",
        "target",
    ]

    def __init__(
        self,
        tp_device: Optional[str] = None,
        tp_device_parameter: Optional[List[str]] = None,
        tp_target: Optional[str] = None,
        tp_target_parameter: Optional[List[str]] = None,
        family: Optional[str] = None,
        firmware: Optional[str] = None,
        prov_firmware: Optional[str] = None,
        audit_log: Optional[str] = None,
        audit_log_key: Optional[str] = None,
        timeout: Optional[int] = None,
        config: Optional[str] = None,
    ) -> None:
        """Initialize the TPHost configuration."""
        config_data = self._get_config_data(config)
        config_dir = os.path.dirname(config) if config else None

        super().__init__(config_data=config_data, config_dir=config_dir)

        if timeout is not None:
            self.config_data["timeout"] = timeout
        if tp_device:
            self.config_data["tp_device"] = tp_device
        if tp_target:
            self.config_data["tp_target"] = tp_target
        if family:
            self.config_data["family"] = family
        if firmware:
            self.config_data["firmware"] = firmware
        if prov_firmware:
            self.config_data["prov_firmware"] = prov_firmware
        if audit_log:
            self.config_data["audit_log"] = audit_log
        if audit_log_key:
            self.config_data["audit_log_key"] = audit_log_key

        self.config_data["tp_device_parameter"] = sanitize_param_struct(
            "tp_device_parameter", tp_device_parameter, self.config_data, self.config_dir
        )
        self.config_data["tp_target_parameter"] = sanitize_param_struct(
            "tp_target_parameter", tp_target_parameter, self.config_data, self.config_dir
        )

        self._validate()

        if "audit_log" not in self.config_data:
            return

        try:
            self.config_data["audit_log"] = find_file(
                self.config_data["audit_log"],
                search_paths=[self.config_dir] if self.config_dir else None,
            )
        except SPSDKError:
            # file doesn't exist yet
            # if config file is defined make the audit_log path relative to it
            if self.config_dir and "audit_log" in self.SCHEMA_MEMBERS:
                self.config_data["audit_log"] = os.path.join(
                    self.config_dir, self.config_data["audit_log"]
                )

    @property
    def tp_target(self) -> str:
        """Trust Provisioning target."""
        return self.config_data["tp_target"]

    @property
    def tp_target_parameter(self) -> Dict[str, Any]:
        """Trust Provisioning target parameters."""
        return self.config_data["tp_target_parameter"]

    @property
    def audit_log(self) -> str:
        """Path to audit log."""
        return self.config_data["audit_log"]

    @property
    def audit_log_key(self) -> str:
        """Path to audit log key."""
        # up until this point 'audit_log_key' was optional
        if "audit_log_key" not in self.config_data:
            raise SPSDKTpError("audit-log-key is not specified")
        return find_file(
            self.config_data["audit_log_key"],
            search_paths=[self.config_dir] if self.config_dir else None,
        )

    @property
    def firmware_data(self) -> Optional[bytes]:
        """OEM firmware content."""
        if not self.config_data.get("firmware"):
            return None
        file_path = find_file(
            self.config_data["firmware"],
            search_paths=[self.config_dir] if self.config_dir else None,
        )
        return load_binary(file_path)

    @property
    def prov_firmware_data(self) -> Optional[bytes]:
        """Provisioning Firmware content."""
        if not self.config_data.get("prov_firmware"):
            return None
        file_path = find_file(
            self.config_data["prov_firmware"],
            search_paths=[self.config_dir] if self.config_dir else None,
        )
        return load_binary(file_path)


class TPConfigConfig(TPBaseConfig):
    """Configuration class for TPConfig app."""

    SCHEMA_MEMBERS = [
        "family",
        "tp_timeout",
        "device",
        "cmpa",
        "cfpa",
        "sb_kek",
        "user_kek",
        "production_quota",
        "oem_log_prk",
        "nxp_prod_cert",
        "nxp_global_attest_cert",
        "oem_id",
    ]

    def __init__(
        self,
        config_file_path: str,
        tp_device: Optional[str] = None,
        tp_device_parameter: Optional[List[str]] = None,
        timeout: Optional[int] = None,
    ) -> None:
        """Initialize the TPConfig configuration."""
        config_dir = os.path.dirname(config_file_path)
        config_data = self._get_config_data(config_file_path)

        super().__init__(config_data=config_data, config_dir=config_dir)

        if timeout is not None:
            self.config_data["timeout"] = timeout
        if tp_device:
            self.config_data["tp_device"] = tp_device
        self.config_data["tp_device_parameter"] = sanitize_param_struct(
            "tp_device_parameter", tp_device_parameter, self.config_data, self.config_dir
        )

        self._validate()


def multiple_tp_dict(multi: Optional[List[str]]) -> Dict[str, str]:
    """Convert even Multiple option to dict in order.

    :param multi: Input List with multiple options.
    :return: Result dictionary
    :raises SPSDKTpError: Problem parsing the input
    """
    try:
        if multi is None:
            return {}
        return dict(item.split("=") for item in multi if isinstance(item, str))
    except Exception as e:
        raise SPSDKTpError(f"Unable to process input {multi}") from e


def sanitize_param_struct(
    param_name: str, cli_params: Optional[List[str]], config_params: dict, config_dir: Optional[str]
) -> dict:
    """Sanitize TP Target/Device parameter settings.

    :param param_name: Name of the parameter group
    :param cli_params: Parameters passed from CLI
    :param config_params: Parameters read from config file
    :param config_dir: Path to config file (for extended search)
    :return: Sanitized parameter config structure
    """
    sanitized_config = config_params.get(param_name, {}) or {}
    sanitized_config.update(multiple_tp_dict(cli_params))
    if "config_file" in sanitized_config:
        config_file_path = sanitized_config["config_file"]
        config_file_path = find_file(
            config_file_path, search_paths=[config_dir] if config_dir else None
        )
        sanitized_config["config_file"] = config_file_path
    return sanitized_config


def print_device_table(intfs: List[TpIntfDescription]) -> str:
    """Prints the List of Interfaces to nice colored table."""
    if len(intfs) == 0:
        return (
            colorama.Fore.RED
            + "Nothing to print - empty interface list!"
            + colorama.Style.RESET_ALL
        )

    # exhaust the chain iterator
    header = list(
        # get all the header elements from interfaces
        # serialize them in to chain: list[list[str]] -> iterator[str]
        itertools.chain.from_iterable(i.as_dict().keys() for i in intfs),
    )
    # remove duplicates whilst preserving the order
    header = list(dict.fromkeys(header))

    table = prettytable.PrettyTable(_sanitize_table_header(header))
    table.align = "l"
    table.header = True
    table.border = True
    table.hrules = prettytable.HEADER
    table.vrules = prettytable.NONE
    for i, intf in enumerate(intfs):
        fields = [
            colorama.Fore.YELLOW + str(i),
            colorama.Fore.MAGENTA + intf.name,
            colorama.Fore.GREEN + intf.description,
        ]

        for field in header[2:]:
            fields.append(colorama.Fore.CYAN + str(intf.as_dict().get(field, "")))  # type: ignore
        table.add_row(fields)

    return table.get_string() + colorama.Style.RESET_ALL


def _sanitize_table_header(header: List[str]) -> List[str]:
    """Sanitize header for use as Table header.

    Capitalize header items.
    Replace "_" with a " " in header items.
    Insert "#" at the first place.
    """
    assert "name" in header
    assert "description" in header
    header = [item.capitalize() for item in header]
    header = [item.replace("_", " ") for item in header]
    header.insert(0, "#")
    return header


def process_tp_inputs(
    tp_type: str,
    tp_parameters: Union[List[str], Dict],
    scan_func: Callable[[Optional[str], Optional[dict]], List[TpIntfDescription]],
    header: Literal["device", "target"],
    print_func: Callable[[str], None],
) -> TpIntfDescription:
    """Process input from config file and/or command line and return corresponding interface.

    In case no interface (device/target) is found, function throws an error.
    In case there are more interfaces found, print out a list of interfaces and exit.

    :param tp_type: Type of tp adapter (scard, blhost, ...)
    :param tp_parameters: Parameters specifying TP Interface
    :param scan_func: Function for scanning TP Interfaces
    :param header: String determining potential error message header (device/target)
    :param print_func: Function for displaying potential error message
    :raises SPSDKTpError: No TP Interface (device/target) found
    :return: Designated Interface Description
    """
    assert tp_type in get_tp_device_types() + get_tp_target_types()
    header_map = {"device": "TP Device", "target": "TP Target"}
    styled_header = header_map[header]
    params = tp_parameters if isinstance(tp_parameters, dict) else multiple_tp_dict(tp_parameters)
    interfaces = scan_func(tp_type, params)
    if not interfaces:
        raise SPSDKTpError(
            f"No {styled_header} found. Search criteria: type={tp_type}; params={params}"
        )
    if len(interfaces) > 1:
        print_func(
            f"{len(interfaces)} {styled_header}s found "
            f"with search criteria type={tp_type}; params={params}"
        )
        print_func(print_device_table(interfaces))
        print_func(f"You need to provide search criteria that fit only one {styled_header}.")
        raise SPSDKTpError()
    return interfaces[0]


###############################################################
# Common options/commands used in both tphost and tpconfig applications


def tp_device_options(options: FC) -> FC:
    """Provides: tp_device: str, tp_device_parameter: List[str]."""
    options = click.option(
        "-dp",
        "--tp-device-parameter",
        type=str,
        multiple=True,
        help="""
The Trusted Device parameters, may be used multiple times to setup TP device interface.
To see available parameters for given TP device, run `tphost/tpconfig device-help`.
Example: -dp id=123456
""",
    )(options)
    options = click.option(
        "-d",
        "--tp-device",
        type=click.Choice(get_tp_device_types(), case_sensitive=False),
        help="The Trusted Device to be used for provisioning target.",
        default=get_tp_device_types()[0] if single_tp_device_adapter() else None,
    )(options)
    return options


def tp_target_options(options: FC) -> FC:
    """Provides: tp_target: str, tp_target_parameter: List[str]."""
    options = click.option(
        "-tp",
        "--tp-target-parameter",
        type=str,
        multiple=True,
        help="""
The Trusted Target parameter, may be used multiple times to setup TP target interface.
To see available parameters for given TP device, run `tphost/tpconfig target-help`.
Example: -tp blhost_baudrate=115200
""",
    )(options)
    options = click.option(
        "-t",
        "--tp-target",
        type=click.Choice(get_tp_target_types(), case_sensitive=False),
        help="The Trusted Target Interface to be used for provisioning process.",
        default=get_tp_target_types()[0] if single_tp_target_adapter() else None,
    )(options)
    return options


@click.command(name="device-help")
@click.option(
    "-d",
    "--tp-device",
    help="Device name to print help, if not used, all devices help will be printed.",
)
def device_help(tp_device: Optional[str] = None) -> None:
    """Command prints help for all devices or optionally only for specified."""
    dev_list = [tp_device] if tp_device else get_tp_device_types()

    for i, dev in enumerate(dev_list):
        click.echo(colorama.Fore.YELLOW + f"#{i}: {dev.upper()} device:")
        click.echo(colorama.Fore.WHITE + get_tp_device_class(dev).get_help() + "\n")
    click.echo(colorama.Style.RESET_ALL)


@click.command(name="target-help")
@click.option(
    "-t",
    "--tp-target",
    help="Target name to print help, if not used, all targets help will be printed.",
)
def target_help(tp_target: Optional[str] = None) -> None:
    """Command prints help for all targets or optionally only for specified."""
    target_list = [tp_target] if tp_target else get_tp_target_types()

    for i, target in enumerate(target_list):
        click.echo(colorama.Fore.YELLOW + f"#{i}: {target.upper()} target:")
        click.echo(colorama.Fore.WHITE + get_tp_target_class(target).get_help() + "\n")
    click.echo(colorama.Style.RESET_ALL)


@click.command(name="list-tptargets", no_args_is_help=not single_tp_target_adapter())
@tp_target_options
def list_tptargets(tp_target: str, tp_target_parameter: List[str]) -> None:
    """Command prints all supported and connected TP targets."""
    tp_targets = scan_tp_targets(tp_target, multiple_tp_dict(tp_target_parameter))
    click.echo(print_device_table(tp_targets))


@click.command(name="list-tpdevices", no_args_is_help=not single_tp_device_adapter())
@tp_device_options
def list_tpdevices(tp_device: str, tp_device_parameter: List[str]) -> None:
    """Command prints all supported and connected TP devices."""
    tp_devices = scan_tp_devices(tp_device, multiple_tp_dict(tp_device_parameter))
    click.echo(print_device_table(tp_devices))


@click.command(name="get-counters", no_args_is_help=False)
@tp_device_options
@click.option(
    "-t",
    "--timeout",
    type=click.IntRange(0, 600, clamp=True),
    help="The target provisioning timeout in seconds.",
)
@click.option(
    "-c",
    "--config",
    type=click.Path(exists=True, dir_okay=False),
    help="Path to configuration file (parameters on CLI take precedence).",
    required=False,
)
def get_counters(
    tp_device: str,
    tp_device_parameter: List[str],
    timeout: int,
    config: str,
) -> None:
    """Get Provisioning counters from TP device."""
    TPHostConfig.SCHEMA_MEMBERS = ["device"]
    tp_config = TPHostConfig(
        tp_device=tp_device,
        tp_device_parameter=tp_device_parameter,
        config=config,
        timeout=timeout,
    )
    device = tp_config.tp_device
    params = tp_config.tp_device_parameter

    tp_interface = process_tp_inputs(
        tp_type=device,
        tp_parameters=params,
        header="device",
        scan_func=scan_tp_devices,
        print_func=click.echo,
    )
    tp_dev = tp_interface.create_interface()
    assert isinstance(tp_dev, TpDevInterface)

    tp_worker = TrustProvisioningHost(
        tpdev=tp_dev,
        tptarget=None,  # type: ignore  # target is not used, we set it to None on purpose
        info_print=click.echo,
    )
    tp_worker.get_counters()
