#!/usr/bin/env python3
# Copyright (c) 2004-present Facebook All rights reserved.
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file.

from typing import Dict, List, Optional, Type, Any, Callable

import importlib.util

from graphql import GraphQLSchema

from .query_parser import (
    ParsedField,
    ParsedObject,
    ParsedOperation,
    ParsedQuery,
    ParsedVariableDefinition,
)
from .utils_codegen import (
    CodeChunk,
    camel_case_to_lower_case,
)

from importlib.abc import Loader

from marshmallow.fields import Field as MarshmallowField

from dataclasses import dataclass

DEFAULT_MAPPING = {
    "ID": "str",
    "String": "str",
    "Int": "int",
    "Float": "Number",
    "Boolean": "bool",
}


@dataclass
class CustomScalar:
    name: str
    type: Type[Any]
    encoder: Optional[Callable[..., Any]] = None
    decoder: Optional[Callable[..., Any]] = None
    mm_field: Optional[MarshmallowField] = None


class DataclassesRenderer:
    def __init__(
        self, schema: GraphQLSchema, config_path: Optional[str] = None
    ) -> None:
        self.schema = schema
        self.config_path = config_path
        self.custom_scalars = {}
        if config_path is not None:
            spec = importlib.util.spec_from_file_location("config", config_path)
            config = importlib.util.module_from_spec(spec)
            assert isinstance(spec.loader, Loader)
            spec.loader.exec_module(config)
            assert hasattr(config, "custom_scalars"), "Custom scalars is not in config"
            self.custom_scalars = getattr(config, "custom_scalars")

    def render(
        self,
        parsed_query: ParsedQuery,
        fragment_name_to_importpath: Dict[str, str],
        enum_name_to_importpath: Dict[str, str],
        input_name_to_importpath: Dict[str, str],
        config_importpath: Optional[str],
    ) -> str:
        buffer = CodeChunk()
        self.__write_file_header(buffer)
        buffer.write("from dataclasses import dataclass, field")
        self.__render_customer_scalars_imports(buffer, config_importpath)
        buffer.write("from gql_client.runtime.variables import encode_variables")
        buffer.write("from gql import gql, Client")
        buffer.write("from gql.transport.exceptions import TransportQueryError")
        buffer.write("from functools import partial")
        buffer.write("from numbers import Number")
        buffer.write(
            "from typing import Any, AsyncGenerator, Dict, List, Generator, Optional"
        )
        buffer.write("from time import perf_counter")
        buffer.write("from dataclasses_json import DataClassJsonMixin, config")
        buffer.write("")
        for fragment_name in sorted(set(parsed_query.used_fragments)):
            importpath = fragment_name_to_importpath[fragment_name]
            buffer.write(
                f"from {importpath} import {fragment_name}, QUERY as {fragment_name}Query"
            )
            buffer.write("")
        enum_names = set()
        for enum in parsed_query.enums:
            enum_names.add(enum.name)
        if enum_names:
            buffer.write(
                "from gql_client.runtime.enum_utils import enum_field_metadata"
            )
            for enum_name in sorted(enum_names):
                importpath = enum_name_to_importpath[enum_name]
                buffer.write(f"from {importpath} import {enum_name}")
            buffer.write("")
        input_object_names = set()
        for input_object in parsed_query.input_objects:
            input_object_names.add(input_object.name)
        if input_object_names:
            for input_object_name in sorted(input_object_names):
                importpath = input_name_to_importpath[input_object_name]
                buffer.write(f"from {importpath} import {input_object_name}")
            buffer.write("")

        sorted_objects = sorted(
            parsed_query.objects,
            key=lambda obj: 1 if isinstance(obj, ParsedOperation) else 0,
        )
        for obj in sorted_objects:
            buffer.write("")
            if isinstance(obj, ParsedObject):
                self.__render_object(parsed_query, buffer, obj)
            elif isinstance(obj, ParsedOperation):
                self.__render_operation(parsed_query, buffer, obj, config_importpath)

        if parsed_query.fragment_objects:
            buffer.write("# fmt: off")
            if parsed_query.used_fragments:
                queries = [
                    f"{fragment_name}Query"
                    for fragment_name in sorted(set(parsed_query.used_fragments))
                ]
                buffer.write(f'QUERY: List[str] = {" + ".join(queries)} + ["""')
            else:
                buffer.write('QUERY: List[str] = ["""')
            buffer.write(parsed_query.query)
            buffer.write('"""]')
            buffer.write("")

        for fragment_obj in parsed_query.fragment_objects:
            self.__render_fragment(parsed_query, buffer, fragment_obj)

        return str(buffer)

    def render_enums(self, parsed_query: ParsedQuery) -> Dict[str, str]:
        result = {}

        for enum in parsed_query.enums + parsed_query.internal_enums:
            buffer = CodeChunk()
            self.__write_file_header(buffer)
            buffer.write("from enum import Enum")
            buffer.write("")
            buffer.write("")
            with buffer.write_block(f"class {enum.name}(Enum):"):
                for value_name, value in enum.values.items():
                    if isinstance(value, str):
                        value = f'"{value}"'

                    buffer.write(f"{value_name} = {value}")
                buffer.write('MISSING_ENUM = ""')
                buffer.write("")
                buffer.write("@classmethod")
                with buffer.write_block(
                    f'def _missing_(cls, value: object) -> "{enum.name}":'
                ):
                    buffer.write("return cls.MISSING_ENUM")
            buffer.write("")
            result[enum.name] = str(buffer)

        return result

    def render_input_objects(
        self, parsed_query: ParsedQuery, config_importpath: Optional[str]
    ) -> Dict[str, str]:
        result = {}

        for input_object in parsed_query.input_objects + parsed_query.internal_inputs:
            buffer = CodeChunk()
            self.__write_file_header(buffer)
            buffer.write("from dataclasses import dataclass, field")
            buffer.write("from functools import partial")
            self.__render_customer_scalars_imports(buffer, config_importpath)
            buffer.write("from numbers import Number")
            buffer.write(
                "from typing import Any, AsyncGenerator, Dict, List, Generator, Optional"
            )
            buffer.write("")
            buffer.write("from dataclasses_json import DataClassJsonMixin, config")
            buffer.write("")
            enum_names = set()
            for enum in input_object.input_enums:
                enum_names.add(enum.name)
            if enum_names:
                buffer.write(
                    "from gql_client.runtime.enum_utils import enum_field_metadata"
                )
                for enum_name in sorted(enum_names):
                    buffer.write(
                        f"from ..enum.{camel_case_to_lower_case(enum_name)} import {enum_name}"
                    )
                buffer.write("")
            input_object_names = set()
            for input_dep in input_object.inputs:
                input_object_names.add(input_dep.name)
            if input_object_names:
                for input_object_name in sorted(input_object_names):
                    buffer.write(
                        f"from ..input.{camel_case_to_lower_case(input_object_name)} "
                        f"import {input_object_name}"
                    )
                buffer.write("")

            buffer.write("")
            self.__render_object(parsed_query, buffer, input_object, True)
            result[input_object.name] = str(buffer)

        return result

    def __render_object(
        self,
        parsed_query: ParsedQuery,
        buffer: CodeChunk,
        obj: ParsedObject,
        is_input: bool = False,
    ) -> None:
        class_parents = (
            "(DataClassJsonMixin)" if not obj.parents else f'({", ".join(obj.parents)})'
        )

        buffer.write("@dataclass(frozen=True)")
        with buffer.write_block(f"class {obj.name}{class_parents}:"):
            # render child objects
            children_names = set()
            for child_object in obj.children:
                if child_object.name not in children_names:
                    self.__render_object(parsed_query, buffer, child_object, is_input)
                children_names.add(child_object.name)

            # render fields
            fields = obj.fields
            if is_input:
                fields = self.__sort_fields(parsed_query, obj.fields)

            for field in fields:
                self.__render_field(parsed_query, buffer, field, is_input)

            # pass if not children or fields
            if not (obj.children or obj.fields):
                buffer.write("pass")

        buffer.write("")

    def __render_fragment(
        self, parsed_query: ParsedQuery, buffer: CodeChunk, obj: ParsedObject
    ) -> None:
        class_parents = (
            "(DataClassJsonMixin)" if not obj.parents else f'({", ".join(obj.parents)})'
        )

        buffer.write("@dataclass(frozen=True)")
        with buffer.write_block(f"class {obj.name}{class_parents}:"):

            # render child objects
            children_names = set()
            for child_object in obj.children:
                if child_object.name not in children_names:
                    self.__render_object(parsed_query, buffer, child_object)
                children_names.add(child_object.name)

            # render fields
            for field in obj.fields:
                self.__render_field(parsed_query, buffer, field)

        buffer.write("")

    def __render_operation(
        self,
        parsed_query: ParsedQuery,
        buffer: CodeChunk,
        parsed_op: ParsedOperation,
        config_importpath: Optional[str],
    ) -> None:
        buffer.write("# fmt: off")
        if len(parsed_query.used_fragments):
            queries = [
                f"{fragment_name}Query"
                for fragment_name in sorted(set(parsed_query.used_fragments))
            ]
            buffer.write(f'QUERY: List[str] = {" + ".join(queries)} + ["""')
        else:
            buffer.write('QUERY: List[str] = ["""')
        buffer.write(parsed_query.query)
        buffer.write('"""')
        buffer.write("]")
        buffer.write("")
        buffer.write("")
        with buffer.write_block(f"class {parsed_op.name}:"):
            # Render children
            for child_object in parsed_op.children:
                self.__render_object(parsed_query, buffer, child_object)

            # Execution functions
            if parsed_op.variables:
                vars_args = ", " + ", ".join(
                    [
                        self.__render_variable_definition(var)
                        for var in parsed_op.variables
                    ]
                )
                variables_dict = (
                    "{"
                    + ", ".join(
                        f'"{var.name}": {var.name}' for var in parsed_op.variables
                    )
                    + "}"
                )
            else:
                vars_args = ""
                variables_dict = "{}"

            assert len(parsed_op.children) == 1
            child = parsed_op.children[0]
            assert len(child.fields) == 1
            query = child.fields[0]
            query_name = query.name
            query_result_type = f"{query.type}"
            if (
                query_result_type not in DEFAULT_MAPPING.keys()
                and query_result_type not in self.custom_scalars.keys()
            ):
                query_result_type = f"{parsed_op.name}Data.{query.type}"
            if query.nullable:
                query_result_type = f"Optional[{query_result_type}]"
            if query.is_list:
                query_result_type = f"List[{query_result_type}]"
            if parsed_op.type in ["query", "mutation"]:
                self.__write_execute_method(
                    buffer,
                    vars_args,
                    query_result_type,
                    variables_dict,
                    parsed_op.name,
                    query_name,
                    config_importpath,
                )
                self.__write_async_execute_method(
                    buffer,
                    vars_args,
                    query_result_type,
                    variables_dict,
                    parsed_op.name,
                    query_name,
                    config_importpath,
                )
            else:
                self.__write_subscribe_method(
                    buffer,
                    vars_args,
                    query_result_type,
                    variables_dict,
                    parsed_op.name,
                    query_name,
                    config_importpath,
                )
                self.__write_async_subscribe_method(
                    buffer,
                    vars_args,
                    query_result_type,
                    variables_dict,
                    parsed_op.name,
                    query_name,
                    config_importpath,
                )

    def __render_customer_scalars_imports(
        self, buffer: CodeChunk, config_importpath: Optional[str]
    ) -> None:
        if not config_importpath:
            return
        scalar_types = set()
        for _, custom_scalar in self.custom_scalars.items():
            if custom_scalar.type.__module__ != "builtins":
                scalar_types.add(custom_scalar.type.__name__)
        types_import_line = ", " + ", ".join(scalar_types) if scalar_types else ""
        buffer.write(
            f"from {config_importpath} import custom_scalars{types_import_line}"
        )

    @staticmethod
    def __write_execute_method(
        buffer: CodeChunk,
        vars_args: str,
        query_result_type: str,
        variables_dict: str,
        operation_name: str,
        query_name: str,
        config_importpath: Optional[str],
    ) -> None:
        buffer.write("# fmt: off")
        buffer.write("@classmethod")
        with buffer.write_block(
            f"def execute(cls, client: Client{vars_args})" f" -> {query_result_type}:"
        ):
            buffer.write(f"variables: Dict[str, Any] = {variables_dict}")
            scalars = ", custom_scalars" if config_importpath else ""
            buffer.write(f"new_variables = encode_variables(variables{scalars})")
            buffer.write("response_text = client.execute(")
            buffer.write('    gql("".join(set(QUERY))), variable_values=new_variables')
            buffer.write(")")
            buffer.write(f"res = cls.{operation_name}Data.from_dict(response_text)")
            buffer.write(f"return res.{query_name}")
        buffer.write("")

    @staticmethod
    def __write_async_execute_method(
        buffer: CodeChunk,
        vars_args: str,
        query_result_type: str,
        variables_dict: str,
        operation_name: str,
        query_name: str,
        config_importpath: Optional[str],
    ) -> None:
        buffer.write("# fmt: off")
        buffer.write("@classmethod")
        with buffer.write_block(
            f"async def execute_async(cls, client: Client{vars_args})"
            f" -> {query_result_type}:"
        ):
            buffer.write(f"variables: Dict[str, Any] = {variables_dict}")
            scalars = ", custom_scalars" if config_importpath else ""
            buffer.write(f"new_variables = encode_variables(variables{scalars})")
            buffer.write("response_text = await client.execute_async(")
            buffer.write('    gql("".join(set(QUERY))), variable_values=new_variables')
            buffer.write(")")
            buffer.write(f"res = cls.{operation_name}Data.from_dict(response_text)")
            buffer.write(f"return res.{query_name}")
        buffer.write("")

    @staticmethod
    def __write_subscribe_method(
        buffer: CodeChunk,
        vars_args: str,
        query_result_type: str,
        variables_dict: str,
        operation_name: str,
        query_name: str,
        config_importpath: Optional[str],
    ) -> None:
        buffer.write("# fmt: off")
        buffer.write("@classmethod")
        with buffer.write_block(
            f"def subscribe(cls, client: Client{vars_args})"
            f" -> Generator[{query_result_type}, None, None]:"
        ):
            buffer.write(f"variables: Dict[str, Any] = {variables_dict}")
            scalars = ", custom_scalars" if config_importpath else ""
            buffer.write(f"new_variables = encode_variables(variables{scalars})")
            buffer.write("subscription = client.subscribe(")
            buffer.write('    gql("".join(set(QUERY))), variable_values=new_variables')
            buffer.write(")")
            with buffer.write_block("for response_text in subscription:"):
                buffer.write(f"res = cls.{operation_name}Data.from_dict(response_text)")
                buffer.write(f"yield res.{query_name}")
        buffer.write("")

    @staticmethod
    def __write_async_subscribe_method(
        buffer: CodeChunk,
        vars_args: str,
        query_result_type: str,
        variables_dict: str,
        operation_name: str,
        query_name: str,
        config_importpath: Optional[str],
    ) -> None:
        buffer.write("# fmt: off")
        buffer.write("@classmethod")
        with buffer.write_block(
            f"async def subscribe_async(cls, client: Client{vars_args})"
            f" -> AsyncGenerator[{query_result_type}, None]:"
        ):
            buffer.write(f"variables: Dict[str, Any] = {variables_dict}")
            scalars = ", custom_scalars" if config_importpath else ""
            buffer.write(f"new_variables = encode_variables(variables{scalars})")
            buffer.write("subscription = client.subscribe_async(")
            buffer.write('    gql("".join(set(QUERY))), variable_values=new_variables')
            buffer.write(")")
            with buffer.write_block("async for response_text in subscription:"):
                buffer.write(f"res = cls.{operation_name}Data.from_dict(response_text)")
                buffer.write(f"yield res.{query_name}")
        buffer.write("")

    @staticmethod
    def __sort_fields(
        parsed_query: ParsedQuery, fields: List[ParsedField]
    ) -> List[ParsedField]:
        def sort_key(field) -> int:
            if field.nullable:
                return 2
            return 0

        return sorted(fields, key=sort_key)

    def __render_variable_definition(self, var: ParsedVariableDefinition):
        var_type = DEFAULT_MAPPING.get(var.type, var.type)

        if var.type in self.custom_scalars.keys():
            var_type = self.custom_scalars[var.type].type.__name__

        if var.is_list:
            return f"{var.name}: List[{var_type}] = []"

        if not var.nullable:
            return f"{var.name}: {var_type}"

        return f'{var.name}: Optional[{var_type}] = {var.default_value or "None"}'

    def __render_field(
        self,
        parsed_query: ParsedQuery,
        buffer: CodeChunk,
        field: ParsedField,
        is_input: bool = False,
    ) -> None:
        enum_names = [e.name for e in parsed_query.enums + parsed_query.internal_enums]
        is_enum = field.type in enum_names
        suffix = ""
        field_type = DEFAULT_MAPPING.get(field.type, field.type)
        if field.is_list:
            field_type = f"List[{field_type}]"

        if is_enum:
            suffix = f" = field(metadata=enum_field_metadata({field_type}))"

        if field_type in self.custom_scalars.keys():
            if (
                self.custom_scalars[field_type].encoder
                or self.custom_scalars[field_type].decoder
                or self.custom_scalars[field_type].mm_field
            ):
                suffix = f' = field(metadata=config(encoder=custom_scalars["{field_type}"].encoder, decoder=custom_scalars["{field_type}"].decoder, mm_field=custom_scalars["{field_type}"].mm_field))'
            field_type = self.custom_scalars[field_type].type.__name__

        if field.nullable:
            if is_input:
                suffix = f" = {field.default_value}"
            buffer.write(f"{field.name}: Optional[{field_type}]{suffix}")
        else:
            buffer.write(f"{field.name}: {field_type}{suffix}")

    @staticmethod
    def __write_file_header(buffer: CodeChunk) -> None:
        buffer.write("#!/usr/bin/env python3")
        buffer.write("# @" + "generated AUTOGENERATED file. Do not Change!")
        buffer.write("")
