from __future__ import annotations

import logging
import subprocess  # nosec
from pathlib import Path
from typing import Union, Dict, Any, List
import asyncio
from urllib.parse import quote_plus

import docker
import json
import os
import time

from yarl import URL

from checkov.common.bridgecrew.vulnerability_scanning.integrations.docker_image_scanning import \
    docker_image_scanning_integration
from checkov.common.bridgecrew.platform_integration import bc_integration
from checkov.common.util.file_utils import decompress_file_gzip_base64
from checkov.common.util.http_utils import request_wrapper, aiohttp_client_session_wrapper
from checkov.common.bridgecrew.platform_key import bridgecrew_dir
from checkov.common.sca.commons import validate_image_id

TWISTCLI_FILE_NAME = 'twistcli'
DOCKER_IMAGE_SCAN_RESULT_FILE_NAME = 'docker-image-scan-results.json'
CHECKOV_SEC_IN_WEEK = 604800
_REDACT_FLAGS = ("--token",)


def redact_token_args(command_args: List[str]) -> List[str]:
    """Return a copy of command_args with the value following any flag in _REDACT_FLAGS replaced by '***'."""
    redacted: List[str] = list(command_args)
    for i, arg in enumerate(redacted[:-1]):
        if arg in _REDACT_FLAGS:
            redacted[i + 1] = "***"
    return redacted


def generate_image_name() -> str:
    return f'repository/image{str(time.time() * 1000)}'


def _get_docker_image_name(docker_image_id: str) -> str:
    try:
        docker_client = docker.from_env()
        image_name: str = docker_client.images.get(docker_image_id).attrs['RepoDigests'][0].split('@')[0]
        return image_name
    except Exception:
        logging.info("Failed to fetch image name.", exc_info=True)
        return generate_image_name()


def _get_dockerfile_content(dockerfile_path: Union[str, "os.PathLike[str]"]) -> str:
    try:
        with open(dockerfile_path) as f:
            return f.read()
    except FileNotFoundError:
        logging.error("Path to Dockerfile is invalid", exc_info=True)
        raise
    except Exception:
        logging.error("Failed to read Dockerfile content", exc_info=True)
        raise


class ImageScanner:
    def __init__(self) -> None:
        self.docker_image_name = ''
        self.dockerfile_content = ''
        self.twistcli_path = Path(bridgecrew_dir) / TWISTCLI_FILE_NAME

    def setup_scan(
            self,
            docker_image_id: str,
            dockerfile_path: Union[str, "os.PathLike[str]"],
            skip_extract_image_name: bool,
    ) -> bool:
        if skip_extract_image_name:
            # Provide a default image name in case the image has not been tagged with a name
            self.docker_image_name = generate_image_name()
        else:
            self.docker_image_name = _get_docker_image_name(docker_image_id)
        self.dockerfile_content = _get_dockerfile_content(dockerfile_path)

        status: bool = self.setup_twistcli()
        return status

    def cleanup_scan(self) -> None:
        if self.twistcli_path.exists():
            os.remove(self.twistcli_path)
            logging.info('twistcli file removed')

    def run_image_scan(self, docker_image_id: str) -> Dict[str, Any]:
        if not validate_image_id(docker_image_id):
            logging.error(f"Invalid docker image id format: {docker_image_id}")
            self.cleanup_scan()
            raise ValueError(f"Invalid docker image id format: {docker_image_id}")

        command_args = [
            str(self.twistcli_path),
            "images", "scan",
            "--address", docker_image_scanning_integration.get_proxy_address(),
            "--token", docker_image_scanning_integration.get_bc_api_key(),
            "--details",
            "--output-file", DOCKER_IMAGE_SCAN_RESULT_FILE_NAME,
            docker_image_id,
        ]
        logging.debug(f"TwistCLI: {redact_token_args(command_args)}")
        try:
            subprocess.run(command_args, check=True, shell=False)  # nosec B603
        except Exception as exc:
            logging.error("Failed to scan image", exc_info=True)
            self.cleanup_scan()
            raise Exception(f"Failed to scan image {docker_image_id}") from exc
        self.cleanup_scan()
        logging.info(f'TwistCLI ran successfully on image {docker_image_id}')

        with open(DOCKER_IMAGE_SCAN_RESULT_FILE_NAME) as docker_image_scan_result_file:
            scan_result: dict[str, Any] = json.load(docker_image_scan_result_file)
        return scan_result

    def scan(self, docker_image_id: str, dockerfile_path: str, skip_extract_image_name: bool = False) -> int:
        try:
            setup_status: bool = self.setup_scan(docker_image_id, dockerfile_path, skip_extract_image_name)
            if not setup_status:
                return 2
            scan_result = self.run_image_scan(docker_image_id)
            exit_code = asyncio.run(
                docker_image_scanning_integration.report_results_async(
                    twistcli_scan_result=scan_result,
                    bc_platform_integration=bc_integration,
                    bc_api_key=docker_image_scanning_integration.get_bc_api_key(),
                    file_path=Path(dockerfile_path),
                    file_content=self.dockerfile_content,
                    docker_image_name=self.docker_image_name,
                )
            )

            logging.info('Docker image scanning results reported to the platform')
            return exit_code
        except Exception:
            logging.error("Failed to scan docker image", exc_info=True)
            return 1

    @staticmethod
    def get_scan_results_from_cache(image_id: str) -> Dict[str, Any] | None:
        try:
            image_id_encode = quote_plus(image_id)
            response = request_wrapper(
                "GET", f"{bc_integration.api_url}/api/v1/vulnerabilities/scan-results/{image_id_encode}",
                headers=bc_integration.get_default_headers("GET"), should_call_raise_for_status=True
            )
            response_json = response.json()
            return ImageScanner._extract_cache_results_for_image(image_id, response_json)

        except Exception:
            logging.debug(
                "Unexpected failure happened during retrieving image scanning result from cache. details are below.\n"
                "Note that the scan is still running. if it is repeated, please report.", exc_info=True)
            return None

    @staticmethod
    async def get_scan_results_from_cache_async(image_id: str) -> Dict[str, Any]:
        """
        This is an async implementation of `get_scan_results_from_cache`. The only change is we're getting a session
        as an input, and the asyncio behavior is managed in the calling method.
        """
        try:
            image_id_encode = quote_plus(image_id)
            url = f"{bc_integration.api_url}/api/v1/vulnerabilities/scan-results/{image_id_encode}"
            headers = bc_integration.get_default_headers("GET")
            logging.debug(f"Invoking API {url}")

            response = await aiohttp_client_session_wrapper("GET", URL(url, encoded=True), headers=headers)
            response_json = await response.json()
            logging.debug(response_json)
            return ImageScanner._extract_cache_results_for_image(image_id, response_json)

        except Exception as e:
            logging.info(f"(IR debug) an Error is raised in (get_scan_results_from_cache_async) for {image_id}.")
            logging.info(
                "Unexpected failure happened during retrieving image scanning result from cache. details are below.\n"
                "Note that the scan is still running. if it is repeated, please report.")
            logging.info(str(e))
            return {}

    @staticmethod
    def _extract_cache_results_for_image(image_id: str, response_json: dict[str, str]) -> dict[str, Any]:
        output_type = response_json.get("outputType")
        logging.info(f"output_type={output_type} returned from cache for image_id={image_id}")
        if output_type == "Result":
            result: dict[str, Any] = json.loads(
                decompress_file_gzip_base64(
                    response_json["outputData"]
                )
            )
            return result
        if output_type is None or output_type == "Error":
            logging.error(response_json.get("outputData"))
            return {}
        logging.info(f"Got an empty result for image={image_id}")
        return {}

    def should_download(self) -> bool:
        if not self.twistcli_path.exists():
            return True
        last_modification = os.stat(self.twistcli_path)
        file_age = (time.time() - last_modification.st_mtime)
        return file_age >= int(os.getenv("CHECKOV_EXPIRATION_TIME_IN_SEC", CHECKOV_SEC_IN_WEEK))

    def setup_twistcli(self) -> bool:
        if self.should_download():
            if not os.path.exists(bridgecrew_dir):
                try:
                    os.makedirs(bridgecrew_dir)
                except FileExistsError:
                    # In multi-processing, this might meet a race condition
                    pass
            self.cleanup_scan()
            status: bool = docker_image_scanning_integration.download_twistcli(image_scanner.twistcli_path)
            return status
        return True


image_scanner = ImageScanner()
