import random
import string
import time

import jwt
import requests
from requests.auth import AuthBase
import urllib3
from tqdm import tqdm
from loguru import logger
from flask import abort
from typing import List, Tuple, Optional

from helpers.time_helper import current_timestamp
from template_pwndoc import TemplatePwndoc, MultilocaleTemplateWrapperPwndoc, TNSTemplatePwndoc
import config
from helpers.file_utils import *


class PwnDocAuth(AuthBase):

    @staticmethod
    def _replace_cookies_with_session_cookies(r: requests.PreparedRequest):
        r.headers.pop('Cookie', None)  # If the request already has a header Cookie than prepare_cookies will not overwrite it.
        r.prepare_cookies(session.cookies)  # Updating cookies on session does not automatically update cookies on prepared request.

    def __call__(self, r):
        if r.path_url.startswith('/api/users'):
            return r

        if session.cookies.get('token') is None:
            _login()
            self._replace_cookies_with_session_cookies(r)
            return r

        jwt_token = session.cookies.get('token', "").replace("JWT%20", "")
        jwt_data = jwt.decode(jwt_token, algorithms=["HS256"], options={"verify_signature": False})
        exp_timestamp = jwt_data.get("exp")

        remaining_seconds = exp_timestamp - current_timestamp()
        if remaining_seconds < 60:  # 60 seconds as a safety margin
            try:
                _refresh_token()
            except AssertionError:
                _login()
            self._replace_cookies_with_session_cookies(r)

        return r


session = requests.Session()
session.auth = PwnDocAuth()
session.verify = not config.PWNDOC_DISABLE_HTTPS_VERIFICATION
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)  # I'm writing one custom warning on config load, the default setting of urllib3 prints it on every request.


def test_connection() -> int:
    resp = session.get(f'{config.PWNDOC_URL}/api/companies', allow_redirects=True)
    return resp.status_code


def _login():
    assert not config.PWNDOC_URL.startswith(config.PWNDOC_NO_CONNECTION_PLACEHOLDER), 'Tried to perform login NO-CONNECTION-TEST domain. Aborting.'
    data = {
        "username": config.PWNDOC_USERNAME,
        "password": config.PWNDOC_PASSWORD,
        "totpToken": "",
    }

    resp = session.post(f"{config.PWNDOC_URL}/api/users/token", data=data)
    assert resp.status_code == 200, f"Login failed. Did you set PwnDoc username ({config.PWNDOC_USERNAME}) and password ({'*'*len(config.PWNDOC_PASSWORD)}) in ENV variables?"


def _refresh_token():
    resp = session.get(f"{config.PWNDOC_URL}/api/users/refreshtoken")
    assert resp.status_code == 200


def upsert_raw_finding(audit_id: str, data: dict):
    finding_id = data.get("_id")
    if finding_id is None:
        url = f"{config.PWNDOC_URL}/api/audits/{audit_id}/findings"
        resp = session.post(url, json=data)
    else:
        url = f"{config.PWNDOC_URL}/api/audits/{audit_id}/findings/{finding_id}"
        resp = session.put(url, json=data)
    assert resp.status_code == 200, f"Upserting of finding failed; URL: {url}; {resp.text}"


def add_templates_for_vulnerability(data: MultilocaleTemplateWrapperPwndoc):
    """ Don't use this directly, rather use PwndocTemplateManager.add_single_locale_template
        This function doesn't fix HTML for PwnDoc, nor does it add it to Importer DB
    """

    url = f"{config.PWNDOC_URL}/api/vulnerabilities"
    fid = data.get_fid()

    pwndoc_json_format = data.get_pwndoc_repr()

    if data._id is None:
        resp = session.post(url, json=[pwndoc_json_format])
    else:
        resp = session.put(f"{url}/{data._id}", json=pwndoc_json_format)

    if resp.status_code == 422 and resp.json()["datas"] == "Vulnerability title already exists":
        logger.warning(f"Rewrite of existing vulnerability '{fid}' is not supported. Skipping.")
    elif resp.status_code in [200, 201]:
        # todo: saving the pwndoc ID mapping here would be nice, though it's not inside the resp
        pass
    else:
        if resp.status_code == 404 and resp.json()["datas"] == "Vulnerability not found":
            FlashLog.error("Attempt to update PwnDoc template failed - the template might have got deleted in meantime, "
                           "or the PwnDoc DB table pwndoc_mapping got desychronized from PwnDoc DB."
                           "Run debug action: Re-synchronize FID to PwnDoc ID mappings.")
        assert False, f"Creation of vulnerability failed: {resp.status_code} {resp.text}"


def filename_from_url_path(url_path: str):
    return f"{url_path.replace('/', '_')}.json"


def get_and_save(url_path: str, filename: Optional[str] = None) -> dict:
    if filename is None:
        filename = filename_from_url_path(url_path)

    url = f"{config.PWNDOC_URL}{url_path}"
    resp = session.get(url)
    if resp.status_code != 200:
        FlashLog.error(f"PwnDoc API failed {url_path} | {resp.status_code} | {resp.text}")
        abort(500)

    def _save_to_file(data: dict, filename: str):
        with open(relative_path(f"api_examples/{filename}"), "w", encoding="utf8") as f:
            json.dump(data, f, indent=4)

    _save_to_file(resp.json(), filename)
    if url_path == "/api/vulnerabilities":
        _save_to_file(resp.json(), filename+f"_backup_{time.time()}.json")
    return resp.json()


def pwndoc_download_json_of_vulnerabilities_and_custom_fields():
    """ Don't use this directly, use PwndocTemplateManager.update_template_db_from_pwndoc instead. """
    get_and_save("/api/vulnerabilities")
    get_and_save("/api/data/custom-fields")


def pwndoc_download_json_of_custom_fields():
    get_and_save("/api/data/custom-fields")


def _download_audits():
    get_and_save("/api/audits")


def refresh_examples(audit_id: Optional[str] = None):
    urls = [
        "/api/vulnerabilities",
        "/api/settings/public",
        "/api/audits",
        "/api/data/languages",
        "/api/data/audit-types",
        "/api/data/custom-fields",
        "/api/data/vulnerability-categories",
        "/api/data/sections",
        "/api/companies",
        "/api/clients",
        "/api/templates",
        "/api/users",
        "/api/users/reviewers",
    ]
    for url in tqdm(urls, desc="Updating API examples"):
        get_and_save(url)

    if audit_id:
        download_audit(audit_id)


def load_audit_file(audit_id: str) -> Optional[dict]:
    """ Note: Use refresh_examples(audit_id) before hand. """
    filepath = relative_path(f"api_examples/{filename_from_url_path('/api/audits/'+audit_id)}")
    data = json_safe_load(filepath)
    return data


def delete_all_templates_in_pwndoc():
    url = f"{config.PWNDOC_URL}/api/vulnerabilities"
    resp = session.delete(url)
    assert resp.status_code == 200, f"Template removal failed: {resp.text}"


def download_audit(audit_id: str):
    urls = [
        f"/api/audits/{audit_id}",
        f"/api/audits/{audit_id}/general",
        f"/api/audits/{audit_id}/network",
    ]
    for url in urls:
        get_and_save(url)


def download_and_return_audit(audit_id: str) -> dict:
    download_audit(audit_id)
    answer = load_audit_file(audit_id)
    if answer is None:
        abort(400, "Audit with this ID doesn't exist.")
    return answer.get('datas', {})


def get_findings_from_audit(audit_id: str) -> List[dict]:
    old_audit = download_and_return_audit(audit_id)

    # note: Findings can be extracted from the Audit endpoint, however if I take the finding from the Audit endpoint
    #   and upsert it using the Finding endpoint, then it breaks custom fields in Pwndoc UI.
    #   Getting the data from Finding endpoint is less efficient, but works.
    findings_from_audit_format = old_audit.get("findings", [])
    findings_from_finding_format = []

    for single_finding in findings_from_audit_format:
        data = download_and_return_finding(audit_id, single_finding.get("_id"))
        findings_from_finding_format.append(data)
    return findings_from_finding_format


def download_and_return_finding(audit_id: str, finding_id: str):
    url_path = f"/api/audits/{audit_id}/findings/{finding_id}"
    return get_and_save(url_path).get("datas", {})


def delete_finding(audit_id: str, finding_id: str):
    url_path = f"/api/audits/{audit_id}/findings/{finding_id}"
    resp = session.delete(f"{config.PWNDOC_URL}{url_path}")
    assert resp.status_code == 200, f"Template removal failed: {resp.text}"


def update_scope(audit_id: str, scope_raw: dict):
    url_path = f"/api/audits/{audit_id}/network"
    resp = session.put(f"{config.PWNDOC_URL}{url_path}", json=scope_raw)

    json_dump(scope_raw, persistent_audit_scopes_file_path(audit_id))

    assert resp.status_code == 200, f"Scope update failed: {resp.text}"


def _get_audits_data() -> List[dict]:
    _download_audits()
    with open(relative_path("api_examples/_api_audits.json"), encoding="utf8") as f:
        audits = json.load(f)
    return audits["datas"]


def get_audits() -> List[Tuple[str, str]]:
    answer = []
    for x in _get_audits_data():
        answer.append((x["_id"], x["name"]))
    return answer


def get_audit(audit_id: str) -> Optional[Dict[str, Any]]:
    for x in _get_audits_data():
        if audit_id == x["_id"]:
            return x
    return None


def does_audit_exist(audit_id: str) -> bool:
    return bool(get_audit(audit_id))


def get_audit_language(audit_id: str) -> Optional[str]:
    audit = get_audit(audit_id)
    return audit["language"] if audit else None


# pwndoc_download_vulnerabilities_and_custom_field()

def _docx_folder() -> str:
    return folder_in_upload_path('pwndoc_generated_report')


def get_docx_filepath(docx_filename: str) -> str:
    return os.path.join(_docx_folder(), docx_filename)


# this is needed because browser has a potentially to short timeout
def download_report(audit_id: str, report_filename: str):
    # note: pass filename, not filepath - filepath can be different between flask and background worker
    report_filepath = get_docx_filepath(report_filename)

    url = f"{config.PWNDOC_URL}/api/audits/{audit_id}/generate"
    resp = session.get(url, timeout=config.RQ_JOB_TIMEOUT)
    assert resp.status_code == 200, f"Report download failed with status code {resp.status_code}"

    with open(report_filepath, "wb") as f:
        f.write(resp.content)


def create_audit(name_section: str = "", locale: str = "cs") -> str:
    data = {
        "name": f"TEST_{name_section}_" + "".join(random.choices(string.ascii_uppercase, k=10)),
        "language": locale,
        "auditType": "Example audit type"
    }
    resp = session.post(f"{config.PWNDOC_URL}/api/audits", data=data)
    assert resp.status_code == 201
    return resp.json()["datas"]["audit"]["_id"]


def get_waiting_template_updates(finding_id: str) -> str:
    # session.get(f"{config.PWNDOC_URL}/api/vulnerabilities/updates/{finding_id}")
    raise NotImplementedError()


def add_template_update(pwndoc_template_dict: dict, pwndoc_id: str, locale: str):
    pwndoc_template_dict['_id'] = pwndoc_id
    pwndoc_template_dict['identifier'] = 1  # sequential ID inside audit
    pwndoc_template_dict['status'] = 1  # 0:done, 1:redacting
    pwndoc_template_dict['paragraphs'] = []  # todo: what is it?

    resp = session.post(f"{config.PWNDOC_URL}/api/vulnerabilities/finding/{locale}", json=pwndoc_template_dict)
    assert resp.status_code == 200


# todo: consider what happens after logout?
# todo: data vs json in requests
