import json
import os
from pathlib import Path
import time
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

from flask import request, redirect, url_for, render_template, flash, Blueprint, abort, send_file
from tqdm import tqdm
from loguru import logger

import config
import pwndoc_api
from template_manager import PwndocTemplateManager
from template_scan2report import TemplateScan2Report
from helpers.flask_multifile_upload import save_multiple_uploaded_files
from helpers.file_utils import json_dump, relative_path, extract_fid_from_filename, zip_folder_recursively, generate_random_folder_path
from helpers.db_models import DbTemplate, DbPwndocMapping
from template_converters import Scan2ReportConvert, PwndocConverter
from template_pwndoc import TNSTemplatePwndoc, MultilocaleTemplateWrapperPwndoc
from helpers.custom_logging import FlashLog
from helpers.flask_response import basic_msg, basic_msg_html

bp = Blueprint('api_templates', __name__, template_folder='templates')


@bp.route("/")
def templates_index():
    return render_template('template_index.html')


@bp.route("/download_from_pwndoc")
@bp.route("/download_from_pwndoc/<output_format>")
@bp.route("/download_from_pwndoc/<output_format>/<locale>")
def download_templates_from_pwndoc_to_scan2report(
        output_format: str = "scan2report_native",
        locale: str = config.LOCALES[0],
        force_locale_folder: Optional[str] = None,
        prefix_title_with_locale: bool = False
):
    """
        This function deletes all local templates and downloads the relevant locale from pwndoc to replace them.
    """
    OUTPUT_FORMAT_MAPPING = {"scan2report_native": True, "scan2report_with_pwndoc_formatting": False}
    PwndocTemplateManager.update_template_db_from_pwndoc()
    TemplateScan2Report.recreate_template_folders_if_necessary()

    if force_locale_folder is None:
        force_locale_folder = locale

    if output_format not in OUTPUT_FORMAT_MAPPING.keys():
        abort(400, description="The selected output format is not supported.")
        return  # this is return is unreachable, but included to silence PyCharm warning

    TemplateScan2Report.delete_template_files(list({locale, "og"}))

    s2r_dicts: List[dict] = PwndocConverter.convert_to_list_of_scan2report_combined_dict(
        output_without_html_tags=OUTPUT_FORMAT_MAPPING[output_format],
        locale=locale
    )
    output_folder = TemplateScan2Report.native_folder_name_for_locale(force_locale_folder)
    for single_template_dict in s2r_dicts:
        output_file_path = os.path.join(output_folder, f"{single_template_dict.get('fid')}.json")
        if Path(output_file_path).is_file():
            logger.warning(f"Template for {single_template_dict.get('fid')} with locale {force_locale_folder} already exists on disk. Skipping to avoid overwriting it.")
            continue

        if prefix_title_with_locale:
            single_template_dict['name'] = f"[{locale.upper()}] {single_template_dict.get('name', '')}"

        json_dump(single_template_dict, output_file_path)

    return basic_msg(f"OK; {len(s2r_dicts)} templates downloaded; Converted to format {output_format} for locale {locale}")


@bp.route("/import_scan2report_to_pwndoc", methods=["GET", "POST"])
def import_scan2report_templates_to_pwndoc():
    if request.method == "GET":
        return render_template("template_import.html", locales=config.LOCALES)

    locale_to_import = request.form.get("template_locale")
    if locale_to_import not in config.LOCALES:
        FlashLog.warning("This locale doesn't exist.")
        return redirect(url_for('api_templates.import_scan2report_templates_to_pwndoc'))

    output_folder = generate_random_folder_path(f"template_import_{locale_to_import}")

    error_msg, saved_filenames = save_multiple_uploaded_files(output_folder, reserved_filenames=set(), allowed_extensions={"json"})
    if len(error_msg) > 0:
        FlashLog.error(error_msg)
        return redirect(url_for('api_templates.import_scan2report_templates_to_pwndoc'))

    saved_filepaths = [os.path.join(output_folder, x) for x in saved_filenames]
    upload_scan2report_templates(locale_to_import, saved_filepaths)
    return redirect(url_for('api_templates.import_scan2report_templates_to_pwndoc'))


def upload_scan2report_templates(locale_to_import: str, saved_filepaths: List[str]):
    # note: requires app or request context
    PwndocTemplateManager.update_template_db_from_pwndoc()  # necessary for mapping new uploads to existing templates
    bundled_templates: Dict[str, Dict[str, TNSTemplatePwndoc]] = defaultdict(dict)

    for single_filepath in tqdm(saved_filepaths, desc="Converting templates"):
        single_filename = os.path.basename(single_filepath)
        fid = extract_fid_from_filename(single_filename)
        with open(single_filepath, encoding="utf8") as f:
            template_json = json.load(f)
        lang_and_templates: Dict[str, TemplateScan2Report] = TemplateScan2Report.parse_dict_for_locale_and_og(fid, locale_to_import, template_json)
        # I don't need to do DB refresh (for pwndoc ID) after every addition because I'm processing only single locale (or locale and og at the same time).
        for locale, single_template in lang_and_templates.items():
            single_tns_template: TNSTemplatePwndoc = Scan2ReportConvert.convert_single_finding_to_tns_pwndoc(single_template, locale)
            assert fid == single_tns_template.fid
            assert locale == single_tns_template.locale
            bundled_templates[fid][locale] = single_tns_template
    del single_filename, fid, single_tns_template, locale, template_json, lang_and_templates, f

    for fid in tqdm(bundled_templates, desc="Uploading templates"):
        dict_of_locales: Dict[str, TNSTemplatePwndoc] = bundled_templates[fid]
        PwndocTemplateManager.add_locales_for_fid(list(dict_of_locales.values()))

    PwndocTemplateManager.update_template_db_from_pwndoc()  # todo: this is workaround for not using PwndocTemplateManager.add_single_locale_template

    FlashLog.info(f"OK; {len(bundled_templates)} templates uploaded")


@bp.route("/export_as_scan2report_native", methods=["GET"])
@bp.route("/export_as_scan2report_native/<output_format>", methods=["GET"])
def export_as_scan2report_native(output_format: str = "scan2report_native"):
    for locale in config.LOCALES:
        download_templates_from_pwndoc_to_scan2report(output_format, locale)
    zip_path = zip_folder_recursively(config.SCAN2REPORT_PLUGINS_FOLDER)
    return send_file(
        zip_path,
        download_name=f"scan2report_templates_{int(time.time())}.zip",
        as_attachment=True
    )


@bp.route("/all_versions/<fid>/<locale>", methods=["GET"])
def count_versioned_template(fid: str, locale: str):
    PwndocTemplateManager.update_template_db_from_pwndoc()
    history = DbTemplate.get_history(fid, locale)
    return [{
        'history_id': x.id,
        'timestamp': x.timestamp,
        'data': json.loads(x.content)
    } for i, x in enumerate(history)]


@bp.route("/version/<history_id>", methods=["GET"])
def show_template_by_history_id(history_id: str):
    try:
        history_id = int(history_id)
    except ValueError:
        return basic_msg("History ID is not a number"), 500

    template = DbTemplate.get(history_id)

    if template is None:
        return basic_msg("Template with this history ID doesn't exist."), 500

    return json.loads(template.content)  # todo: maybe also return time?


@bp.route("/version_to_pwndoc/<int:history_id>", methods=["GET"])
def upload_version_to_pwndoc_for_comparison(history_id: int):
    PwndocTemplateManager.update_template_db_from_pwndoc()

    db_template = DbTemplate.get(history_id)
    og_content: TNSTemplatePwndoc = db_template.decode_content()

    cur_db_template = DbTemplate.get_template(og_content.fid, og_content.locale)
    cur_content: TNSTemplatePwndoc = cur_db_template.decode_content()

    pwndoc_id = DbPwndocMapping.get_pwndoc_id_from_fid(og_content.fid)
    pwndoc_api.add_template_update(og_content.to_pwndoc_dict(), pwndoc_id, db_template.locale)

    return basic_msg_html(f'Template revision {history_id} for FID "{og_content.fid}", locale "{og_content.locale}" was uploaded to PwnDoc. '
                          f"Find the template '{cur_content.title}' in PwnDoc. It's current state is unchanged, proposed update is the revision. "
                          f'<a href="/vulnerabilities">PwnDoc templates</a>')


@bp.route("/current/<fid>/<locale>", methods=["GET"])
def show_current_version_of_template(fid: str, locale: str):
    template = DbTemplate.get_template(fid, locale)
    return show_template_by_history_id(str(template.id if template else -1))


@bp.route("/compare_versions", methods=["GET"])
def compare_versions(fid: str = None, locale: str = None):
    PwndocTemplateManager.update_template_db_from_pwndoc()
    all_latest_versions = DbTemplate.get_all_current_templates()
    fids = set([x.fid for x in all_latest_versions])

    # history: List[Tuple[str, str]] = []  # history_id, "history_id: date and time"

    return render_template('template_history_single.html', fids=fids, locales=config.LOCALES
                           # , chosen_fid=fid, chosen_locale=locale, history=history
    )


@bp.route("/update_db", methods=["GET"])
@bp.route("/update_db/<pwndoc_template_id>", methods=["GET"])
def update_db(pwndoc_template_id: str = ""):
    PwndocTemplateManager.update_template_db_from_pwndoc()  # todo: we don't have to update everything, single template would be enough
    return basic_msg("OK")


@bp.route("/history", methods=["GET"])
def whole_history():
    PwndocTemplateManager.update_template_db_from_pwndoc()
    all_templates = DbTemplate.get_all_history()
    return render_template('template_history_chronological.html', all_templates=all_templates)


class ProblemFinders:
    @staticmethod
    def find_template_locales_without_fid(current_multilocale_templates_in_pwndoc: List[MultilocaleTemplateWrapperPwndoc]) -> List[Dict[str, str]]:
        templates_without_fids: List[Dict[str, str]] = []
        for single_multilocale_finding in current_multilocale_templates_in_pwndoc:
            single_multilocale_finding: TNSTemplatePwndoc = single_multilocale_finding
            for single_finding in single_multilocale_finding.details:

                if single_finding.fid.strip() == "":
                    templates_without_fids.append({
                        'pwndoc_id': MultilocaleTemplateWrapperPwndoc._id,
                        'locale': single_finding.locale,
                        'title': single_finding.title
                    })
        return templates_without_fids

    @staticmethod
    def get_scan2report_dicts(all_templates: Dict[str, Dict[str, TNSTemplatePwndoc]], include_og: bool = True) -> Dict[str, Dict[str, dict]]:
        scan2report_templates_dicts: Dict[str, Dict[str, dict]] = {}

        for fid in all_templates:
            scan2report_templates_dicts[fid] = {}
            for locale, tns_template in all_templates[fid].items():
                if locale == 'og' and include_og is False:  # OG can have anything, only proper locales should be synchronized
                    continue
                scan2report_templates_dicts[fid][locale] = PwndocConverter.convert_to_single_scan2report_dict(tns_template)
        return scan2report_templates_dicts

    @staticmethod
    def get_shared_fields() -> List[str]:
        return ['severity', 'attack_complexicity', 'system_impact', 'links', 'author', 'ignore_pluginoutput']

    @classmethod
    def find_templates_with_non_synchronized_shared_fields(cls, scan2report_templates_dicts: Dict[str, Dict[str, dict]]) -> List[Tuple[str, str]]:
        templates_with_non_shared_fields: List[Tuple[str, str]] = []
        shared_fields = cls.get_shared_fields()
        for fid in scan2report_templates_dicts:
            for single_field_key in shared_fields:
                vals = []
                for _, scan2report_dict in scan2report_templates_dicts[fid].items():
                    vals.append(scan2report_dict[single_field_key])
                all_same = len(vals) == 0 or all(vals[0] == elem for elem in vals)
                if not all_same:
                    templates_with_non_shared_fields.append((fid, single_field_key))

        return templates_with_non_shared_fields


@bp.route("/warnings", methods=["GET"])
def warnings():
    pwndoc_api.pwndoc_download_json_of_vulnerabilities_and_custom_fields()

    current_multilocale_templates_in_pwndoc: List[MultilocaleTemplateWrapperPwndoc] = PwndocTemplateManager._parse_pwndoc_multilang_templates()
    templates_without_fids = ProblemFinders.find_template_locales_without_fid(current_multilocale_templates_in_pwndoc)

    all_templates: Dict[str, Dict[str, TNSTemplatePwndoc]] = PwndocTemplateManager.get_all_templates_as_dict_of_locales()
    scan2report_templates_dicts: Dict[str, Dict[str, dict]] = ProblemFinders.get_scan2report_dicts(all_templates, include_og=False)

    templates_with_non_shared_fields: List[Tuple[str, str]] = ProblemFinders.find_templates_with_non_synchronized_shared_fields(scan2report_templates_dicts)
    shared_fields: List[str] = ProblemFinders.get_shared_fields()

    return render_template('template_warnings.html',
                           templates_without_fids=templates_without_fids,
                           templates_with_non_shared_fields=templates_with_non_shared_fields,
                           shared_fields=shared_fields
    )
