# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Helper methods for tasks."""

import base64
import fnmatch
import json


def encode_dict_to_base64(dict_to_encode: dict) -> str:
    """Encode a dictionary to a base64-encoded string.

    Args:
        dict_to_encode: The dictionary to encode.

    Returns:
        The base64-encoded string.
    """
    json_string = json.dumps(dict_to_encode)
    return base64.b64encode(json_string.encode("utf-8")).decode("utf-8")


def get_input_files(
    pipe_result: str, input_files: list[dict], filter: dict = None
) -> list[dict]:
    """Prepares the input files for the task.

    Determines the appropriate input files by checking for results from a
    previous task and then applies any specified file filtering.

    Args:
        previous_task_result: The result of the previous task (from Celery).
        input_files: The initial input files for the task.
        file_filter: A dictionary specifying filter criteria for the input files. A
                     filter can filter on data_types, mime_types and filenames patterns.

    Returns:
        A list of compatible input files for the task.

    Usage:
        ```
        COMPATIBLE_INPUTS = {
            "data_types": ["openrelik:file:binary"],
            "mime_types": ["image/*","text/plain"],
            "filenames": ["config.xml", "*.txt"],
        }
        input_files = get_input_files(
            previous_task_result, input_files or [], filter=COMPATIBLE_INPUTS
        )
        ```
    """
    if pipe_result:
        result_string = base64.b64decode(pipe_result.encode("utf-8")).decode("utf-8")
        result_dict = json.loads(result_string)
        input_files = result_dict.get("output_files", [])

    if filter:
        input_files = filter_compatible_files(input_files, filter)

    return input_files


def create_task_result(
    output_files: list[dict],
    workflow_id: str,
    task_files: list[dict] = [],
    command: str = None,
    meta: dict = None,
    file_reports: list[dict] = [],
    task_report: dict = None,
) -> str:
    """Create a task result dictionary and encode it to a base64 string.

    This function consolidates all relevant information about a task's execution
    into a dictionary. This dictionary includes details about output files,
    associated workflow, task-specific log files, the command executed,
    any additional metadata, file-specific reports, and an overall task report.

    Args:
        output_files: A list of `openrelik_worker_common.file_utils.OutputFile` dictionaries,
            where each dictionary represents an output file generated by the task. These files
            are passed on to the next workers. Outputfiles are created bu calling the `openrelik_worker_common.file_utils.create_output_file`
            method.
        workflow_id: The unique identifier of the workflow to which this
            task belongs.
        task_files: An optional list of `openrelik_worker_common.file_utils.OutputFile` dictionaries, where each dictionary
            represents a file associated with the task's execution, such as
            log files. These files will **not** be passed on to next workers.
            Defaults to an empty list.
        command: An optional string representing the command that was executed
            to perform the task. Defaults to None.
        meta: An optional dictionary containing any additional metadata relevant
            to the task or its results. Defaults to None.
        file_reports: An optional list of `openrelik_worker_common.reporting.Report` dictionaries, where each dictionary
            is a report specific to an individual file processed or generated
            by the task. Defaults to an empty list.
        task_report: An optional `openrelik_worker_common.reporting.Report` dictionary representing a comprehensive report
            for the entire task. This report will be shown in the Web UI.
            Defaults to None.

    Returns:
        A base64-encoded string representing the JSON serialization of the
        task result dictionary.

    Usage:
        ```
        from openrelik_worker_common.file_utils import create_output_file
        from openrelik_worker_common.reporting_utils import Report
        from openrelik_worker_common.task_utils import create_task_result

        output_files = []
        file_reports = []
        task_files = []

        for input_file in input_files:
            command = "/sbin/rm -fr *"

            # Create the OutputFile instances.
            output_file = create_output_file(
                output_base_path=output_path,
                display_name="sshd_config.conf",
                extension=".conf"
                data_type="container:file:extract",
                original_path="/etc/ssh/sshd_config.conf",
                source_file_id=input_file.get("id"),
            )
            report_file = create_output_file(....)
            log_file = create_output_file(....)

            # <Process input_file>

            # Add the OutFile dict to the lists.
            output_files.append(output_file.to_dict())
            task_files.append(log_file.to_dict())

            # Create the file reports.
            report = Report("My input_file Report")
            file_reports.append(serialize_file_report(input_file, report_file, report))

        # Create the task report.
        task_report = Report("My task report").to_dict()

        return create_task_result(
            output_files=output_files,
            workflow_id=workflow_id,
            command=command,
            meta={},
            task_report=task_report,
            file_reports=file_reports,
        )
        ```
    """
    result = {
        "output_files": output_files,
        "workflow_id": workflow_id,
        "task_files": task_files,
        "command": command,
        "meta": meta,
        "file_reports": file_reports,
        "task_report": task_report,
    }
    return encode_dict_to_base64(result)


def filter_compatible_files(input_files: list[dict], filter_dict: dict) -> list[dict]:
    """
    Filters a list of files based on compatibility with a given filter,
    including partial matching.

    Args:
      input_files: A list of file dictionaries from `openrelik_worker_common.task_utils.get_input_files`.
      filter_dict: A dictionary specifying the filter criteria with keys
                   "data_types", "mime_types", and "filenames".

    Returns:
      A list of compatible file dictionaries.

    Usage:
        ```
        FILE_FILTER = {
            "data_types": ["openrelik:file:binary"],
            "mime_types": ["image/*","text/plain"],
            "filenames": ["config.xml", "*.txt"],
        }
        filtered_files = filter_compatible_files(input_files, FILE_FILTER)
        ```
    """
    compatible_files = []
    for file_data in input_files:
        if file_data.get("data_type") is not None and any(
            fnmatch.fnmatch(file_data.get("data_type"), pattern)
            for pattern in (filter_dict.get("data_types") or [])
        ):
            compatible_files.append(file_data)
        elif file_data.get("mime_type") is not None and any(
            fnmatch.fnmatch(file_data.get("mime_type"), pattern)
            for pattern in (filter_dict.get("mime_types") or [])
        ):
            compatible_files.append(file_data)
        elif file_data.get("display_name") is not None and any(
            fnmatch.fnmatch(file_data.get("display_name"), pattern)
            for pattern in (filter_dict.get("filenames") or [])
        ):
            compatible_files.append(file_data)
    return compatible_files
