import io
import json
import pprint
import re
import textwrap
from typing import Any

from mitmproxy import http
from mitmproxy.utils import strutils


def curl_command(flow: http.HTTPFlow) -> str:
    data = "curl "

    request = flow.request.copy()
    request.decode(strict=False)

    for k, v in request.headers.items(multi=True):
        data += "-H '%s:%s' " % (k, v)

    if request.method != "GET":
        data += "-X %s " % request.method

    data += "'%s'" % request.url

    if request.content:
        data += " --data-binary '%s'" % strutils.bytes_to_escaped_str(
            request.content, escape_single_quotes=True
        )

    return data


def python_arg(arg: str, val: Any) -> str:
    if not val:
        return ""
    if arg:
        arg += "="
    arg_str = "{}{},\n".format(arg, pprint.pformat(val, 79 - len(arg)))
    return textwrap.indent(arg_str, " " * 4)


def python_code(flow: http.HTTPFlow):
    code = io.StringIO()

    def writearg(arg, val):
        code.write(python_arg(arg, val))

    code.write("import requests\n")
    code.write("\n")
    if flow.request.method.lower() in ("get", "post", "put", "head", "delete", "patch"):
        code.write("response = requests.{}(\n".format(flow.request.method.lower()))
    else:
        code.write("response = requests.request(\n")
        writearg("", flow.request.method)
    url_without_query = flow.request.url.split("?", 1)[0]
    writearg("", url_without_query)

    writearg("params", list(flow.request.query.fields))

    headers = flow.request.headers.copy()
    # requests adds those by default.
    for x in (":authority", "host", "content-length"):
        headers.pop(x, None)
    writearg("headers", dict(headers))
    try:
        if "json" not in flow.request.headers.get("content-type", ""):
            raise ValueError()
        writearg("json", json.loads(flow.request.text))
    except ValueError:
        writearg("data", flow.request.content)

    code.seek(code.tell() - 2)  # remove last comma
    code.write("\n)\n")
    code.write("\n")
    code.write("print(response.text)")

    return code.getvalue()


def locust_code(flow):
    code = textwrap.dedent(
        """	
        from locust import HttpLocust, TaskSet, task	
        class UserBehavior(TaskSet):	
            def on_start(self):	
                ''' on_start is called when a Locust start before any task is scheduled '''	
                self.{name}()	
            @task()	
            def {name}(self):	
                url = self.locust.host + '{path}'	
                {headers}{params}{data}	
                self.response = self.client.request(	
                    method='{method}',	
                    url=url,{args}	
                )	
            ### Additional tasks can go here ###	
        class WebsiteUser(HttpLocust):	
            task_set = UserBehavior	
            min_wait = 1000	
            max_wait = 3000	
"""
    ).strip()

    name = re.sub("\W|^(?=\d)", "_", flow.request.path.strip("/").split("?", 1)[0])
    if not name:
        new_name = "_".join([str(flow.request.host), str(flow.request.timestamp_start)])
        name = re.sub("\W|^(?=\d)", "_", new_name)

    path_without_query = flow.request.path.split("?")[0]

    args = ""
    headers = ""

    def conv(x):
        return strutils.bytes_to_escaped_str(x, escape_single_quotes=True)

    if flow.request.headers:
        lines = [
            (conv(k), conv(v))
            for k, v in flow.request.headers.fields
            if conv(k).lower() not in [":authority", "host", "cookie"]
        ]
        lines = ["            '%s': '%s',\n" % (k, v) for k, v in lines]
        headers += "\n        headers = {\n%s        }\n" % "".join(lines)
        args += "\n            headers=headers,"

    params = ""
    if flow.request.query:
        lines = [
            "            %s: %s,\n" % (repr(k), repr(v))
            for k, v in flow.request.query.collect()
        ]
        params = "\n        params = {\n%s        }\n" % "".join(lines)
        args += "\n            params=params,"

    data = ""
    if flow.request.content:
        data = "\n        data = '''%s'''\n" % conv(flow.request.content)
        args += "\n            data=data,"

    code = code.format(
        name=name,
        path=path_without_query,
        headers=headers,
        params=params,
        data=data,
        method=flow.request.method,
        args=args,
    )

    return code


def locust_task(flow):
    code = locust_code(flow)
    start_task = len(code.split("@task")[0]) - 4
    end_task = -19 - len(code.split("### Additional")[1])
    task_code = code[start_task:end_task]

    return task_code


def url(flow):
    return flow.request.url


EXPORTERS = [
    ("content", "c", None),
    ("headers+content", "h", None),
    ("url", "u", url),
    ("as curl command", "r", curl_command),
    ("as python code", "p", python_code),
    ("as locust code", "l", locust_code),
    ("as locust task", "t", locust_task),
]