import json
from dataclasses import asdict
from importlib import import_module
from io import StringIO
from operator import attrgetter
from textwrap import indent
from typing import Any, Callable, ClassVar, Dict, List, Optional, Protocol, Tuple, Type

from . import factories, queries
from .steps import (
    AssertDelete,
    AssertInsert,
    AssertResponseJson,
    AssertStatusCode,
    AssertTemplateUsed,
    AssertUpdate,
    Code,
    CodeComment,
    DjangoTestClient,
    EmptyLine,
    EndTestClass,
    EndTestFunction,
    EndTestMethod,
    EndTimeTravel,
    FactoryCreate,
    Import,
    ModelCreate,
    ModelUpdate,
    ModuleImports,
    RegisterMocket,
    StartTimeTravel,
    Step,
    TestClass,
    TestFunction,
    TestMethod,
)


def json_encoder(obj):
    if isinstance(obj, set):
        return list(obj)
    if isinstance(obj, Import):
        return {"import": obj.render()}
    raise TypeError(
        f"Object of type {obj.__class__.__name__} is not JSON serializable"
    )  # pragma: no cover


class Plan:
    interesting_steps = (
        AssertInsert,
        AssertResponseJson,
        AssertStatusCode,
        AssertTemplateUsed,
        AssertUpdate,
        DjangoTestClient,
        FactoryCreate,
        ModelCreate,
        ModelUpdate,
        RegisterMocket,
    )

    def __init__(self, steps, pytest, context):
        self.steps = list(steps)
        self.pytest = pytest
        self.context = context

    def render(self):
        contents = StringIO()

        imports = set()
        for step in self.steps:
            imports.update(step.get_imports(self.pytest))

        indent_count = 0
        for step in self.steps:
            indent_prefix = "    " * indent_count
            if isinstance(step, ModuleImports):
                for import_step in sorted(imports, key=attrgetter("import_path")):
                    contents.write(indent(import_step.render(), indent_prefix))

            contents.write(indent(step.render(self.pytest), indent_prefix))
            indent_count += step.indent_delta

        return contents.getvalue()

    def as_json(self):
        steps = [
            (step.__class__.__name__, asdict(step))
            for step in self.steps
            if isinstance(step, self.interesting_steps)
        ]
        return json.dumps(steps, default=json_encoder)


def header_comments(context):
    if context["now"]:
        yield CodeComment(f"Auto generated by Kolo at {context['now']}.")
    else:
        yield CodeComment("Auto generated by Kolo.")
    yield CodeComment(f"Kolo version: {context['kolo_version']}")
    if context["django_version"]:
        yield CodeComment(f"Django version: {context['django_version']}")
    for trace_id in context["_traces"]:
        yield CodeComment(f"Trace ID: {trace_id}")


def start_time_travel(section, context):
    if "request_timestamp" in section:
        request_timestamp = repr(section["request_timestamp"])
        if context["time_travel_tick"]:
            args: Tuple[str, ...] = (request_timestamp, "tick=True")
        else:
            args = (request_timestamp,)
        yield StartTimeTravel(
            context["time_travel_call"],
            args,
            imports={Import(context["time_travel_import"])},
        )


def end_time_travel(section):
    if "request_timestamp" in section:
        yield EndTimeTravel()


def call_test_client(section, pytest):
    if "response" in section:
        request = section["request"]
        post_data = request["post_data"]
        post_data.pop("csrfmiddlewaretoken", None)
        yield DjangoTestClient(
            request["method"].lower(),
            request["path_info"],
            request["query_params"],
            post_data,
            section["request_headers"],
        )
        yield EmptyLine()
        yield AssertStatusCode(section["response"]["status_code"])
        if section["response_json"]:
            yield AssertResponseJson(section["response_json"])
        for template_name in section["template_names"]:
            yield AssertTemplateUsed(template_name)
        yield EmptyLine()


def build_steps(context, pytest):
    yield from header_comments(context)
    yield EmptyLine()
    yield ModuleImports()
    yield EmptyLine()
    yield EmptyLine()
    if not pytest:
        yield TestClass(
            context["test_class"],
            ("TestCase",),
            {Import("from django.test import TestCase")},
        )
    yield Code("@mocketize(strict_mode=True)", {Import("from mocket import mocketize")})
    if pytest:
        yield Code("@pytest.mark.django_db()", {Import("import pytest")})
        yield TestFunction(context["test_name"], ("client",))
    else:
        yield TestMethod(context["test_name"])

    for section in context["sections"]:
        yield from start_time_travel(section, context)

        for create_fixture in section["sql_fixtures"]:
            if isinstance(create_fixture, queries.DjangoCreate):
                yield ModelCreate.from_fixture(create_fixture)
            elif isinstance(create_fixture, queries.DjangoUpdate):
                yield ModelUpdate.from_fixture(create_fixture)
            elif isinstance(create_fixture, factories.Factory):  # pragma: no branch
                yield FactoryCreate.from_fixture(create_fixture)
        if len(section["sql_fixtures"]) > 0:
            yield EmptyLine()

        for outbound_request in section["outbound_request_frames"]:
            yield RegisterMocket.from_outbound_request(outbound_request)
        if len(section["outbound_request_frames"]) > 0:
            yield EmptyLine()

        yield from call_test_client(section, pytest)

        for fixture in section["asserts"]:
            if isinstance(fixture, queries.AssertInsert):
                yield AssertInsert.from_fixture(fixture)
            elif isinstance(fixture, queries.AssertUpdate):
                yield AssertUpdate.from_fixture(fixture)
            elif isinstance(fixture, queries.AssertDelete):  # pragma: no branch
                yield AssertDelete.from_fixture(fixture)

        yield from end_time_travel(section)

    if pytest:
        yield EndTestFunction()
    else:
        yield EndTestMethod()
        yield EndTestClass()


def run_plan_hooks(plan, hooks):
    for hook in hooks:
        plan = hook(plan)
    return plan


def run_step_hooks(plan, hooks):
    for step_class, step_hooks in hooks.items():
        indexed_steps = step_class.get_steps(plan.steps)
        new_steps = []
        for index_slice, steps in indexed_steps:
            if len(steps) == 1:
                steps = steps[0]
            for hook in step_hooks:
                steps = hook(steps)
            new_steps.append((index_slice, steps))
        for index_slice, steps in reversed(new_steps):
            plan.steps[index_slice] = steps
    return plan


PlanHook = Callable[[Plan], Plan]
StepHook = Callable[[List[Step]], List[Step]]
plan_hooks: List[PlanHook] = []
step_hooks: Dict[Type[Step], List[StepHook]] = {}


def load_hooks(config):
    imports = config.get("test_generation", {}).get("hook_imports", [])
    for import_path in imports:
        import_module(import_path)  # pragma: no cover
    return plan_hooks, step_hooks


def plan_hook(func: PlanHook):
    plan_hooks.append(func)
    return func


def step_hook(step_class: Type[Step]):
    def decorator(func: StepHook):
        step_hooks.setdefault(step_class, []).append(func)
        return func

    return decorator
