import ast
import jinja2

from collections.abc import Mapping
from datetime import date, datetime, timedelta
from enum import Enum
from typing import Dict
from uuid import UUID

from clipped.utils.serialization import (
    date_serialize,
    datetime_serialize,
    timedelta_serialize,
    uuid_serialize,
)

from polyaxon.exceptions import PolyaxonSchemaError
from polyaxon.polyaxonfile.specs.libs.engine import get_engine
from polyaxon.polyaxonfile.specs.sections import Sections
from polyaxon.polyflow import ParamSpec

try:
    import numpy as np
except (ImportError, ModuleNotFoundError):
    np = None


class PolyaxonfileParser:
    """Parses the Polyaxonfile."""

    engine = get_engine()

    @staticmethod
    def _get_section_data(section_data):
        if hasattr(section_data, "to_dict"):
            return section_data.to_dict()
        return section_data

    @classmethod
    def _get_section(cls, config, section):
        if not hasattr(config, section):
            return None
        section_data = getattr(config, section)
        if isinstance(section_data, list):
            return [cls._get_section_data(d) for d in section_data]

        return cls._get_section_data(section_data)

    @staticmethod
    def get_parsed_params(param_spec: Dict[str, ParamSpec]) -> Dict:
        parsed_params = {"params": {}}
        as_arg = []
        for param in param_spec or {}:
            parsed_params[param] = param_spec[param].get_display_value()
            # Add params level
            if param_spec[param].is_requested:
                as_arg.append(param_spec[param].as_arg())
                parsed_params["params"][param] = param_spec[param].to_parsed_param()

        parsed_params["params"]["as_args"] = [a for a in as_arg if a]
        return parsed_params

    @classmethod
    def parse_operation(
        cls, config, param_spec: Dict[str, ParamSpec]
    ):  # pylint:disable=too-many-branches
        parsed_params = cls.get_parsed_params(param_spec)

        parsed_data = {Sections.VERSION: config.version, Sections.KIND: config.kind}

        if config.name:
            parsed_data[Sections.NAME] = config.name
        if config.description:
            parsed_data[Sections.DESCRIPTION] = config.description
        if config.tags:
            parsed_data[Sections.TAGS] = config.tags
        contexts = getattr(config, Sections.CONTEXTS)
        if contexts:
            parsed_data[Sections.CONTEXTS] = [
                cls.parse_io(io.to_dict(), parsed_params) for io in contexts
            ]
        inputs = getattr(config, Sections.INPUTS)
        if inputs:
            parsed_data[Sections.INPUTS] = [
                cls.parse_io(io.to_dict(), parsed_params) for io in inputs
            ]
        outputs = getattr(config, Sections.OUTPUTS)
        if outputs:
            parsed_data[Sections.OUTPUTS] = [
                cls.parse_io(io.to_dict(), parsed_params) for io in outputs
            ]

        # Check workflow
        for section in Sections.PARSING_SECTIONS:
            config_section = cls._get_section(config, section)
            if config_section:
                parsed_data[section] = cls.parse_expression(
                    config_section, parsed_params
                )

        parsed_data = cls._parse_conditions(parsed_data)

        config_section = cls._get_section(config, Sections.RUN)
        if config_section:
            parsed_data[Sections.RUN] = config_section
        config_section = cls._get_section(config, Sections.BUILD)
        if config_section:
            parsed_data[Sections.BUILD] = config_section
        config_section = cls._get_section(config, Sections.HOOKS)
        if config_section:
            parsed_data[Sections.HOOKS] = config_section

        return parsed_data

    @classmethod
    def _parse_conditions(cls, parsed_data: Dict) -> Dict:
        if Sections.CONDITIONS in parsed_data and not isinstance(
            parsed_data[Sections.CONDITIONS], str
        ):
            parsed_data[Sections.CONDITIONS] = str(parsed_data[Sections.CONDITIONS])
        if Sections.SKIP_ON_UPSTREAM_SKIP in parsed_data and not isinstance(
            parsed_data[Sections.SKIP_ON_UPSTREAM_SKIP], str
        ):
            parsed_data[Sections.SKIP_ON_UPSTREAM_SKIP] = str(
                parsed_data[Sections.SKIP_ON_UPSTREAM_SKIP]
            )
        if Sections.IS_APPROVED in parsed_data and not isinstance(
            parsed_data[Sections.IS_APPROVED], str
        ):
            parsed_data[Sections.IS_APPROVED] = str(parsed_data[Sections.IS_APPROVED])
        return parsed_data

    @classmethod
    def _parse_hook(cls, config, param_spec: Dict[str, ParamSpec]):
        hook = cls.parse_section(config, param_spec=param_spec, parse_params=True)
        return cls._parse_conditions(hook)

    @classmethod
    def parse_hooks(cls, config, param_spec: Dict[str, ParamSpec]):
        hooks = config.hooks or []
        return [
            cls._parse_hook(hook.to_dict(), param_spec=param_spec) for hook in hooks
        ]

    @classmethod
    def parse_runtime(cls, parsed_data, param_spec: Dict[str, ParamSpec]):
        config_section = cls.parse_section(
            parsed_data.get(Sections.RUN), param_spec=param_spec, parse_params=True
        )
        if config_section:
            parsed_data[Sections.RUN] = config_section
        return parsed_data

    @classmethod
    def parse_distributed_runtime(
        cls, parsed_data, param_spec: Dict[str, Dict[str, ParamSpec]]
    ):
        run_section = parsed_data.get(Sections.RUN)
        config_section = {}
        for k in run_section:
            if k in param_spec:
                config_section[k] = cls.parse_section(
                    run_section[k], param_spec=param_spec[k], parse_params=True
                )
            else:
                config_section[k] = run_section[k]
        if config_section:
            parsed_data[Sections.RUN] = config_section
        return parsed_data

    @classmethod
    def parse_section(
        cls, config_section, param_spec: Dict[str, ParamSpec], parse_params: bool = True
    ):
        if parse_params:
            param_spec = cls.get_parsed_params(param_spec)
        if config_section:
            return cls.parse_expression(config_section, param_spec)
        return config_section

    @classmethod
    def parse_io(
        cls,
        expression: Dict,
        params: Dict[str, ParamSpec],
        check_operators: bool = False,
    ):
        # Arg format should not be converted
        arg_format = expression.pop("argFormat", None)
        value = cls.parse_expression(expression, params, check_operators)
        if arg_format:
            value["argFormat"] = arg_format
        return value

    @classmethod
    def parse_expression(  # pylint:disable=too-many-branches
        cls, expression, params: Dict, check_operators: bool = False
    ):
        try:
            return cls._parse_expression(expression, params, check_operators)
        except jinja2.exceptions.TemplateError as e:
            raise PolyaxonSchemaError(
                "Encountered a problem parsing the template, "
                "please make sure your variables are resolvable. "
                "Error: {}".format(repr(e))
            )

    @classmethod
    def _parse_expression(  # pylint:disable=too-many-branches
        cls, expression, params: Dict, check_operators: bool = False
    ):
        if isinstance(expression, (int, float, complex, type(None))):
            return expression
        if isinstance(expression, Enum):
            return expression.value
        if isinstance(expression, datetime):
            return datetime_serialize(expression)
        if isinstance(expression, date):
            return date_serialize(expression)
        if isinstance(expression, timedelta):
            return timedelta_serialize(expression)
        if isinstance(expression, UUID):
            return uuid_serialize(expression)
        if np and isinstance(expression, np.integer):
            return int(expression)
        if np and isinstance(expression, np.floating):
            return float(expression)
        if isinstance(expression, Mapping):
            if len(expression) == 1:
                old_key, value = list(expression.items())[0]
                # always parse the keys, they must be base object or evaluate to base objects
                key = cls._parse_expression(old_key, params)
                if check_operators and cls.is_operator(key):
                    return cls._parse_operator({key: value}, params)
                else:
                    return {key: cls.parse_expression(value, params, check_operators)}

            new_expression = {}
            for k, v in expression.items():
                new_expression.update(
                    cls._parse_expression({k: v}, params, check_operators)
                )
            return new_expression

        if isinstance(expression, list):
            return list(
                cls.parse_expression(v, params, check_operators) for v in expression
            )
        if isinstance(expression, tuple):
            return tuple(
                cls.parse_expression(v, params, check_operators) for v in expression
            )
        if isinstance(expression, str):
            return cls._evaluate_expression(expression, params, check_operators)

    @classmethod
    def _evaluate_expression(cls, expression, params, check_operators):
        try:
            result = cls.engine.from_string(expression).render(**params)
        except (ValueError, TypeError) as e:
            raise PolyaxonSchemaError(
                "Encountered a problem parsing the template, "
                "please make sure your variables are resolvable. "
                "Error: {}".format(repr(e))
            )
        if result == expression:
            try:
                return ast.literal_eval(result)
            except (ValueError, SyntaxError):
                pass
            return result
        return cls.parse_expression(result, params, check_operators)

    @classmethod
    def _parse_operator(cls, expression, params):
        k, v = list(expression.items())[0]
        op = Sections.OPERATORS[k].from_dict(v)
        return op.parse(parser=cls, params=params)

    @staticmethod
    def is_operator(key):
        return key in Sections.OPERATORS
