import base64
import glob
import os
import random
import shutil

import requests
from loguru import logger
import time
import json
from typing import Dict, Optional, List

import api_templates
import config
from helpers.file_utils import relative_path
from pathlib import Path

from pwndoc_api import session, test_connection


class InitialData:
    __NEW_MAPPING_FILEPATH = relative_path(f'user_data/pwndoc-init/new_id_mapping.json')

    def __init__(self):
        self.main_thread: bool = False

    def run_init(self) -> bool:
        self.setup_first_user()  # this can set self.main_thread
        # login happens automatically
        self.upload_initial_data()
        logger.info("PwnDoc Init completed")
        return self.main_thread

    def setup_first_user(self):
        if not self._is_db_clean(max_seconds=30):  # checks if DB is empty, wait if necessary
            logger.info("DB seems to be (at least partially) initialized - skipping initialization.")
            return

        self.main_thread = self._add_pwndoc_user(first_user=True)

        # The thread that created the user will continue. Other threads will block until user is created (or timeout).
        if not self.main_thread:
            self.wait_until_pwndoc_user_is_ready(wait_max_x_seconds=30)
            return

    def upload_initial_data(self):
        self._upload_universal('_api_templates.json')
        self._upload_universal('_api_data_languages.json')
        self._upload_universal('_api_data_sections.json')
        self._upload_universal('_api_data_vulnerability-categories.json')
        self._upload_universal('_api_data_audit-types.json')  # requires: templates, languages, sections
        self._upload_universal('_api_data_custom-fields.json')  # requires: everything

    @staticmethod
    def wait_until_pwndoc_user_is_ready(wait_max_x_seconds=30):
        while wait_max_x_seconds > 0 and test_connection() != 200:
            time.sleep(1)
            wait_max_x_seconds -= 1
        assert test_connection() == 200, "After 30 seconds everything should be ready, but isn't. Restarting thread."

    @staticmethod
    def _is_db_clean(max_seconds: float = 30) -> bool:
        for i in range(10):
            try:
                resp = session.get(f"{config.PWNDOC_URL}/api/users/init")
                if resp.status_code == 200 and resp.json().get("datas") is True:
                    logger.info("Clean instance detected. Forcing import to DB.")
                    return True
                return False
            except requests.exceptions.ConnectionError:
                if config.APP_INIT_PWNDOC_DISALLOW_WAITING:
                    break
                logger.warning(f"Connection to {config.PWNDOC_URL} failed. Waiting before retry.")
                time.sleep(max_seconds/10 + random.uniform(-0.1, 0.1))

        logger.warning(f"Connection to {config.PWNDOC_URL} failed. Not retrying.")
        return False

    @staticmethod
    def _add_pwndoc_user(username: str = config.PWNDOC_USERNAME, password: str = config.PWNDOC_PASSWORD, firstname: str = "Pwndoc",
                         lastname: str = "Importer", first_user: bool = False) -> bool:

        init_data = {
            "username": username,
            "password": password,
            "firstname": firstname,
            "lastname": lastname
        }
        url = f"{config.PWNDOC_URL}/api/users"

        if first_user:
            url += "/init"
        else:
            init_data["role"] = "admin"

        resp = session.post(url, data=init_data)
        if resp.status_code == 403 and resp.json()["datas"] == "Already Initialized":
            logger.info("Init failed - already initialized, possibly by another thread. Skipping.")
            return False
        elif resp.status_code == 422 and resp.json()["datas"] == "Username already exists":
            logger.info(f"User '{username}' already exists. Skipping.")
            return False
        assert resp.status_code == 201, f"Init failed: {resp.text}"
        return True

    def _save_new_id_mapping(self, old_id: str, new_id: str):
        try:
            with open(self.__NEW_MAPPING_FILEPATH, encoding="utf8") as f:
                content = json.load(f)
        except:
            content = {}

        content[old_id] = new_id
        with open(self.__NEW_MAPPING_FILEPATH, "w", encoding="utf8") as f:
            json.dump(content, f, indent=4)

    def _load_new_id_mapping(self, old_id: str) -> str:
        for i in range(10):
            try:
                with open(self.__NEW_MAPPING_FILEPATH, encoding="utf8") as f:
                    content = json.load(f)
                return content.get(old_id, old_id)
            except json.decoder.JSONDecodeError:
                logger.info("Loading for new ID mapping failed. Sleeping before retry.")
                time.sleep(random.uniform(0, 1))

        assert False, "Loading for new ID mapping failed - even on final retry. Killing."

    @staticmethod
    def _get_api_path_from_filename(filename: str) -> str:
        url_path = filename.replace("_", "/").rsplit(".", 1)[0]
        return f"{config.PWNDOC_URL}{url_path}"

    @staticmethod
    def _get_data_from_init_file(filename: str) -> List[dict]:
        with open(relative_path(f'user_data/pwndoc-init/{filename}'), encoding="utf8") as f:
            content = json.load(f)["datas"]
        return content

    def _upload_universal(self, filename: str):
        content = self._get_data_from_init_file(filename)

        for single_item in content:

            if filename == "_api_templates.json":
                single_item['file'] = self._load_template_file(single_item['name'] + "." + single_item['ext']).decode()

            if filename == "_api_data_audit-types.json":
                for locale_dict in single_item.get("templates", []):
                    locale_dict["template"] = self._load_new_id_mapping(locale_dict["template"])

            self._create_single_item(self._get_api_path_from_filename(filename), single_item)

    def _create_single_item(self, url: str, orig_data: dict):
        resp = session.post(url, json=orig_data)
        if resp.status_code == 201:
            if orig_data.get("_id"):
                self._save_new_id_mapping(orig_data["_id"], resp.json()["datas"]["_id"])
        else:
            if not resp.json().get("data", "").endswith("already exists"):
                logger.info(f"Skipping upload duplicate: {resp.text}")
            else:
                logger.warning(f"Upload failed: {resp.text}. Skipping.")

    @staticmethod
    def _load_template_file(filename: str) -> bytes:
        template_filepath = f"user_data/pwndoc-init/{filename}"
        with open(relative_path(template_filepath), "rb") as f:
            template_file = f.read()
            template_file_b64 = base64.b64encode(template_file)
            return template_file_b64

    # -- later updates

    def find_custom_field_by_label(self, datas: List[dict], label: str) -> Optional[dict]:
        one_field = list(filter(lambda x: x.get('label') == label, datas))
        return one_field[0] if len(one_field) else None

    def insert_vuln_custom_field_if_doesnt_exist(self, label: str):
        CUSTOM_FIELDS_FILENAME = '_api_data_custom-fields.json'

        file_datas = self._get_data_from_init_file(CUSTOM_FIELDS_FILENAME)
        local_field = self.find_custom_field_by_label(file_datas, label)
        if local_field is None:
            logger.warning(f"Invoked addition of new custom field {label}, however no field with such label is in {CUSTOM_FIELDS_FILENAME}. Skipping.")
            return

        # login is outside in PwnDocUpdate
        url = self._get_api_path_from_filename(CUSTOM_FIELDS_FILENAME)
        resp = session.get(url)
        assert resp.status_code == 200, f'Python Requests auth mechanism failed for PwnDoc failed? {resp.status_code}'
        pwndoc_datas = resp.json()["datas"]
        pwndoc_field = self.find_custom_field_by_label(pwndoc_datas, label)

        if pwndoc_field is None:
            logger.info(f"Adding Vulnerability custom field '{label}' to PwnDoc.")
            self._create_single_item(url, local_field)


class PwnDocUpdate:
    @staticmethod
    def add_new_fields_if_needed():
        _id = InitialData()
        _id.insert_vuln_custom_field_if_doesnt_exist('ignore_pluginoutput_checkbox')

    @staticmethod
    def setup_scan2report_plugin_folder():
        # todo: this is Scan2Report only, move it outside of PwnDocUpdate class
        if Path(config.SCAN2REPORT_ALTERNATIVE_FOLDER).exists():
            return
        # Path(config.TEMPLATE_FOLDER_SCAN2REPORT).mkdir(parents=True, exist_ok=False)
        assert os.path.join(config.SCAN2REPORT_ALTERNATIVE_FOLDER, 'plugins') == config.SCAN2REPORT_PLUGINS_FOLDER

        shutil.copytree(os.path.join(config.SCAN2REPORT_ORIGINAL_FOLDER, 'plugins'), os.path.join(config.SCAN2REPORT_ALTERNATIVE_FOLDER, 'plugins'))
        shutil.copytree(os.path.join(config.SCAN2REPORT_ORIGINAL_FOLDER, 'templates'), os.path.join(config.SCAN2REPORT_ALTERNATIVE_FOLDER, 'templates'))

    @staticmethod
    def upload_templates_from_scan2report_repository():
        # This requires app context because it needs access to sqlite db.
        for locale in config.LOCALES:
            if locale == 'og':
                return
            try:
                template_filepaths = glob.glob(os.path.join(config.SCAN2REPORT_ORIGINAL_FOLDER, 'plugins', locale, '*.json'))
                api_templates.upload_scan2report_templates(locale, template_filepaths)
            except Exception as e:
                logger.warning(f"Upload of initial templates failed for locale {locale}. Skipping. Reason: {e}")
                raise