import textwrap
import typing
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Set, Tuple, TypeVar

import networkx
import protogen

from py_gen_ml.extensions_pb2 import CLI
from py_gen_ml.logging.setup_logger import setup_logger
from py_gen_ml.plugin.common import (
    WRAP_WIDTH,
    generate_docstring,
    get_element_subgraphs,
    get_extension_value,
    py_import_for_source_file_derived_file,
    snake_case,
)
from py_gen_ml.plugin.constants import (
    BASE_MODEL_ALIAS,
    CLI_ARGS_ALIAS,
    PGML_ALIAS,
    SWEEP_MODEL_ALIAS,
)
from py_gen_ml.plugin.generator import Generator
from py_gen_ml.typing.some import some

logger = setup_logger(__name__)

T = TypeVar('T')


@dataclass
class FieldPath:
    """A path to a field in a message."""
    path: List[str]
    field: protogen.Field


class CliArgsGenerator(Generator):
    """
    A generator that creates CLI argument models based on the input protobuf files.

    The generated file will be imported by the py-gen-ml CLI to generate CLI argument models for the given protobuf files.
    """

    def __init__(self, gen: protogen.Plugin, suffix: str) -> None:
        """
        Initialize the CliArgsGenerator.

        Args:
            gen (protogen.Plugin): The protogen plugin instance.
            suffix (str): The suffix to be added to the generated file names.

        This class is responsible for generating CLI argument models based on
        the input protobuf files. It creates pydantic models that can be used
        with typer for command-line interfaces.
        """
        super().__init__(gen)
        self._suffix = suffix

    def _generate_code_for_file(self, file: protogen.File) -> None:
        """
        Generate the args models corresponding to a file.

        Args:
            file (protogen.File): The input file.

        Returns:
            protogen.GeneratedFile: The generated file.
        """
        if not any(
            (cli := get_extension_value(message, 'cli', CLI)) is not None and cli.enable for message in file.messages
        ):
            return
        g = self._gen.new_generated_file(
            file.proto.name.replace('.proto', self._suffix),
            file.py_import_path,
        )
        g.P('# Autogenerated code. DO NOT EDIT.')
        g.P(f'import py_gen_ml as {PGML_ALIAS}')
        g.P('import typing')
        g.P()
        g.P('import pydantic')
        g.P('import typer')
        g.P()

        import_file_path = str(Path(file.proto.name.replace('.proto', '_base')).name)
        g.P(f'from . import {import_file_path} as {BASE_MODEL_ALIAS}')
        g.P()
        g.P()

        dependency_subgraphs = get_element_subgraphs(file)
        for subgraph in dependency_subgraphs:
            for message in networkx.topological_sort(subgraph):
                if not isinstance(message, protogen.Message):
                    continue
                self._generate_message_args(g, message, subgraph)
                self._generate_entrypoint_script_for_message(message)

        self._run_yapf(g)

    def _generate_message_args(
        self,
        g: protogen.GeneratedFile,
        message: protogen.Message,
        graph: networkx.MultiDiGraph,
    ) -> None:
        """
        Generate the 'Args' counterpart of a given message.

        Args:
            g (protogen.GeneratedFile): The generated file to which we append the generated Python code.
            message (protogen.Message): The proto message that we are currently processing.
            graph (networkx.DiGraph): The networkx graph that contains the message.
        """
        cli_proto = get_extension_value(message, 'cli', CLI)
        if cli_proto is None or not cli_proto.enable:
            return

        g.P(f'class {message.proto.name}Args({PGML_ALIAS}.YamlBaseModel):')
        g.set_indent(4)

        generate_docstring(g, message)

        explicit_paths: Set[Tuple[str, ...]] = set()
        for arg in cli_proto.arg:
            path = arg.path.split('.')
            field = self._get_field_by_path(message, path)
            g.P(f'{arg.name}: {self.field_to_annotation(field, path)}')
            generate_docstring(g, field)
            explicit_paths.add(tuple(path))

        descendant_fields = self._gather_descendant_fields(message, graph, explicit_paths)
        for name, field_paths in descendant_fields.items():
            if len(field_paths) != 1:
                continue
            field_path = field_paths[0]
            g.P(f'{name}: {self.field_to_annotation(field=field_path.field, path=field_path.path)}')
            generate_docstring(g, field_path.field)

        g.P()
        g.P()
        g.set_indent(0)

    def _generate_entrypoint_script_for_message(self, message: protogen.Message) -> None:
        """
        Generate the entrypoint script for a given message.

        Args:
            message (protogen.Message): The message to generate the entrypoint script for.
        """
        cli_proto = get_extension_value(message, 'cli', CLI)
        if cli_proto is None or not cli_proto.enable:
            return
        g = self._gen.new_generated_file(
            f'{snake_case(message.proto.name)}_entrypoint.py',
            message.parent_file.py_import_path,
        )
        sweep_py_import = py_import_for_source_file_derived_file(g, '_sweep')
        sweep_py_import = self._prepend_python_import(sweep_py_import)

        base_py_import = py_import_for_source_file_derived_file(g, '_base')
        base_py_import = self._prepend_python_import(base_py_import)

        cli_args_py_import = py_import_for_source_file_derived_file(g, '_cli_args')
        cli_args_py_import = self._prepend_python_import(cli_args_py_import)

        g.P(f'import {base_py_import} as {BASE_MODEL_ALIAS}')
        g.P(f'import {sweep_py_import} as {SWEEP_MODEL_ALIAS}')
        g.P(f'import {cli_args_py_import} as {CLI_ARGS_ALIAS}')
        g.P(f'import typer')
        g.P(f'import py_gen_ml as {PGML_ALIAS}')
        g.P(f'import optuna')
        g.P('import typing')
        g.P()
        g.P(f'app = typer.Typer(pretty_exceptions_enable=False)')
        g.P()
        g.P(f'def run_trial(')
        g.set_indent(4)
        g.P(f'{snake_case(message.proto.name)}: {BASE_MODEL_ALIAS}.{message.proto.name},')
        g.P('trial: typing.Optional[optuna.Trial] = None')
        g.set_indent(0)
        g.P(') -> typing.Union[float, typing.Sequence[float]]:')
        g.set_indent(4)
        g.P(f"\"\"\"")
        lines = textwrap.wrap(
            f'Run a trial with the given values for {snake_case(message.proto.name)}. The sampled '
            f'hyperparameters have already been added to the trial.',
            width=WRAP_WIDTH - 4,
        )
        for line in lines:
            g.P(line)
        g.P("\"\"\"")
        g.P('# TODO: Implement this function')
        g.P('return 0.0')
        g.set_indent(0)
        g.P()
        g.P(f'@{PGML_ALIAS}.pgml_cmd(app=app)')
        g.P(f'def main(')
        g.set_indent(4)
        g.P(f'config_paths: typing.List[str] = typer.Option(..., help="Paths to config files"),')
        g.P(f'sweep_paths: typing.List[str] = typer.Option(')
        g.set_indent(8)
        g.P('default_factory=list,')
        g.P('help="Paths to sweep files"')
        g.set_indent(4)
        g.P('),')
        g.P(f'cli_args: {CLI_ARGS_ALIAS}.{message.proto.name}Args = typer.Option(...),')
        g.set_indent(0)
        g.P(f') -> None:')
        g.set_indent(4)
        g.P(f'{snake_case(message.proto.name)} = {BASE_MODEL_ALIAS}.{message.proto.name}.from_yaml_files(config_paths)')
        g.P(f'{snake_case(message.proto.name)} = {snake_case(message.proto.name)}.apply_cli_args(cli_args)')
        g.P(f'if len(sweep_paths) == 0:')
        g.set_indent(8)
        g.P(f'run_trial({snake_case(message.proto.name)})')
        g.P('return')
        g.set_indent(4)
        g.P(
            f'{snake_case(message.proto.name)}_sweep = {SWEEP_MODEL_ALIAS}.{message.proto.name}Sweep.from_yaml_files(sweep_paths)',
        )
        g.P()
        g.P('def objective(trial: optuna.Trial) -> typing.Union[')
        g.set_indent(8)
        g.P('float,')
        g.P('typing.Sequence[float]')
        g.set_indent(4)
        g.P(']:')
        g.set_indent(8)
        g.P(f'optuna_sampler = {PGML_ALIAS}.OptunaSampler(trial)')
        g.P(f'{snake_case(message.proto.name)}_patch = optuna_sampler.sample({snake_case(message.proto.name)}_sweep)')
        g.P(
            f'{snake_case(message.proto.name)}_patched = {snake_case(message.proto.name)}.merge({snake_case(message.proto.name)}_patch)',
        )
        g.P(f'objective_value = run_trial({snake_case(message.proto.name)}_patched, trial)')
        g.P('return objective_value')
        g.set_indent(4)
        g.P()
        g.P("study = optuna.create_study(direction='maximize')")
        g.P('study.optimize(objective, n_trials=100)')
        g.set_indent(0)
        g.P()
        g.P()
        g.P('if __name__ == "__main__":')
        g.set_indent(4)
        g.P('app()')

    def _gather_descendant_fields(
        self,
        message: protogen.Message,
        graph: networkx.MultiDiGraph,
        explicit_paths: Set[Tuple[str, ...]],
    ) -> Dict[str, list[FieldPath]]:
        """
        Gathers descendant fields of the given message.

        It goes down the dependency hierarchy to find field names at leaf nodes.

        Args:
            message (protogen.Message): The message of which to find descendant
                fields.
            graph (networkx.MultiDiGraph): A networkx graph that contains the message.

        Returns:
            OrderedDict[str, List[FieldPath]]: A mapping of field name to a list FieldPath objects.
        """
        fields = OrderedDict[str, List[FieldPath]]()
        for descendant in [message, *networkx.ancestors(graph, message)]:
            if isinstance(descendant, protogen.Enum):
                continue
            raw_paths = typing.cast(
                List[Tuple[Any, Any, str]],
                list(
                    networkx.all_simple_edge_paths(graph, source=descendant, target=message),
                ),
            )
            # Filter out paths that are already in explicit_paths
            paths = [tuple(edge[2] for edge in reversed(edges)) for edges in raw_paths]
            if len(paths) == 0:

                for field in descendant.fields:
                    if field.kind == protogen.Kind.MESSAGE:
                        continue

                    name = typing.cast(str, field.py_name)  # type: ignore
                    if (name,) in explicit_paths:
                        continue

                    if name in fields:
                        fields[name].append(FieldPath(path=[name], field=field))
                    else:
                        fields[name] = [FieldPath(path=[name], field=field)]
                continue
            elif len(paths) == 1:
                shortcut_paths = [()]
            else:
                max_len = max(len(path) for path in paths)
                for shortcut_len in range(1, max_len + 1):
                    shortcut_paths: list[tuple[str, ...]] = [edges[-shortcut_len:] for edges in paths]
                    if len(set(shortcut_paths)) == len(paths):
                        break
                else:
                    raise ValueError(f'Could not find a unique path for message {message.proto.name}')

            for field in descendant.fields:
                if field.kind == protogen.Kind.MESSAGE:
                    continue
                for shortcut_path, path in zip(shortcut_paths, paths):
                    field_py_name = typing.cast(str, field.py_name)  # type: ignore
                    if (*path, field_py_name) in explicit_paths:
                        continue
                    name = '_'.join([*shortcut_path, field_py_name])
                    if name in fields:
                        fields[name].append(FieldPath(path=[*path, field_py_name], field=field))
                    else:
                        fields[name] = [FieldPath(path=[*path, field_py_name], field=field)]
        return fields

    def field_to_annotation(self, field: protogen.Field, path: List[str]) -> str:
        """
        Get the annotation for a field and a path.

        The path is what the field maps to in the base model.

        Args:
            field (protogen.Field): Field in the protobuf Message
            path (List[str]): Path to the field in the base model that this field
                maps to.

        Returns:
            List[str]: A list of strings to append.
        """
        annotation = self.field_to_type(field)
        if field.is_list() and len(path) == 0:
            annotation = f'List[{annotation}]'
        typer_option = ''

        if len(path) > 0:
            path_str = '.'.join(path)
            if field.location.leading_comments == '':
                typer_option = f", typer.Option(help=\"Maps to {path_str}\")"
            else:
                typer_option = f", typer.Option(help=\"{field.location.leading_comments.rstrip('.')}. Maps to '{path_str}'\")"
            annotation = f"typing.Annotated[typing.Optional[{annotation}]{typer_option}, pydantic.Field(None), {PGML_ALIAS}.ArgRef(\"{path_str}\"),]"
        else:
            if field.location.leading_comments == '':
                typer_option = f", typer.Option(help=\"Maps to {field.py_name}\")"
            else:
                typer_option = f", typer.Option(help=\"{field.location.leading_comments.rstrip('.')}. Maps to '{field.py_name}'\")"
            annotation = f'typing.Annotated[typing.Optional[{annotation}]{typer_option}, pydantic.Field(None),]'
        return annotation

    def _get_field_by_path(self, message: protogen.Message, path: List[str]) -> protogen.Field:
        """
        Get the field at a given path in a message.

        Args:
            message (protogen.Message): The message to get the field from.
            path (List[str]): The path to the field.

        Returns:
            protogen.Field: The field at the given path.

        Raises:
            ValueError: The field at the given path was not found.
        """
        for field in message.fields:
            if field.py_name == path[0]:
                if len(path) == 1:
                    return field
                if field.kind == protogen.Kind.MESSAGE:
                    return self._get_field_by_path(some(field.message), path[1:])
        else:
            raise ValueError(f'Field {path[0]} not found in message {message.proto.name}')

    def field_to_type(self, field: protogen.Field) -> str:
        """
        Converts a protogen field to a type.

        Args:
            field (protogen.Field): Field in the proto message.

        Raises:
            ValueError: The field kind is not supported.

        Returns:
            str: A string representation of the type.
        """
        if field.kind == protogen.Kind.MESSAGE:
            return f'{some(field.message).py_ident.py_name}Args'
        elif field.kind == protogen.Kind.DOUBLE:
            return 'float'
        elif field.kind == protogen.Kind.FLOAT:
            return 'float'
        elif field.kind == protogen.Kind.INT64:
            return 'int'
        elif field.kind == protogen.Kind.UINT64:
            return 'int'
        elif field.kind == protogen.Kind.INT32:
            return 'int'
        elif field.kind == protogen.Kind.FIXED64:
            return 'int'
        elif field.kind == protogen.Kind.FIXED32:
            return 'int'
        elif field.kind == protogen.Kind.BOOL:
            return 'bool'
        elif field.kind == protogen.Kind.STRING:
            return 'str'
        elif field.kind == protogen.Kind.BYTES:
            return 'bytes'
        elif field.kind == protogen.Kind.UINT32:
            return 'int'
        elif field.kind == protogen.Kind.ENUM:
            return f'{BASE_MODEL_ALIAS}.{some(field.enum).py_ident.py_name}'
        elif field.kind == protogen.Kind.SFIXED32:
            return 'int'
        elif field.kind == protogen.Kind.SFIXED64:
            return 'int'
        elif field.kind == protogen.Kind.SINT32:
            return 'int'
        elif field.kind == protogen.Kind.SINT64:
            return 'int'
        else:
            raise ValueError(f'Unknown field kind: {field.kind}')
