from __future__ import annotations

import os
from collections import defaultdict
from pathlib import Path
from typing import Dict, List, Literal, TypedDict, Union

from ruamel.yaml import YAML
from ruamel.yaml.comments import CommentedMap
from typing_extensions import NotRequired

from metricflow.dataflow.sql_table import SqlTable
from metricflow.inference.models import InferenceResult, InferenceSignalType
from metricflow.inference.renderer.base import InferenceRenderer

yaml = YAML()


class RenderedTimeColumnConfigTypeParams(TypedDict):  # noqa: D
    time_granularity: Literal["day"]


class RenderedColumnConfig(TypedDict):  # noqa: D
    name: str
    type: str
    type_params: NotRequired[RenderedTimeColumnConfigTypeParams]


class ConfigFileRenderer(InferenceRenderer):
    """Writes inference results to a set of config files."""

    UNKNOWN_FIELD_VALUE = "FIXME"

    def __init__(self, dir_path: Union[str, Path], overwrite: bool = False) -> None:
        """Initializes the renderer.

        dir_path: The path to the config directory
        overwrite: If set to False, will raise error if the directory exists
        """
        dir_path = os.path.abspath(dir_path)

        if not overwrite and os.path.exists(dir_path):
            raise ValueError("ConfigFileRender.overwrite is False but path exists.")

        if os.path.isfile(dir_path):
            raise ValueError("ConfigFileRenderer `dir_path` is a file.")

        self.dir_path = dir_path

    def _get_filename_for_table(self, table: SqlTable) -> str:
        return os.path.abspath(os.path.join(self.dir_path, f"{table.sql}.yaml"))

    def _fixme(self, comment: str) -> str:
        return f"FIXME: {comment}"

    def _render_entity_columns(self, results: List[InferenceResult]) -> List[CommentedMap]:
        type_map = {
            InferenceSignalType.ID.PRIMARY: "primary",
            InferenceSignalType.ID.FOREIGN: "foreign",
            InferenceSignalType.ID.UNIQUE: "unique",
        }

        rendered: List[CommentedMap] = [
            CommentedMap(
                {
                    "name": result.column.column_name,
                    "type": type_map.get(result.type_node, ConfigFileRenderer.UNKNOWN_FIELD_VALUE),
                }
            )
            for result in results
            if result.type_node.is_subtype_of(InferenceSignalType.ID.UNKNOWN)
        ]

        return rendered

    def _render_dimension_columns(self, results: List[InferenceResult]) -> List[CommentedMap]:
        type_map = {
            InferenceSignalType.DIMENSION.TIME: "time",
            InferenceSignalType.DIMENSION.PRIMARY_TIME: "time",
            InferenceSignalType.DIMENSION.CATEGORICAL: "categorical",
        }

        rendered: List[CommentedMap] = []
        for result in results:
            if not result.type_node.is_subtype_of(InferenceSignalType.DIMENSION.UNKNOWN):
                continue

            result_data: CommentedMap = CommentedMap(
                {
                    "name": result.column.column_name,
                    "type": type_map.get(result.type_node, ConfigFileRenderer.UNKNOWN_FIELD_VALUE),
                }
            )

            if result_data["type"] == ConfigFileRenderer.UNKNOWN_FIELD_VALUE:
                result_data.yaml_add_eol_comment(self._fixme("unknown field value"), "type")

            if result.type_node.is_subtype_of(InferenceSignalType.DIMENSION.TIME):
                type_params: CommentedMap = CommentedMap({"time_granularity": "day"})
                if result.type_node.is_subtype_of(InferenceSignalType.DIMENSION.PRIMARY_TIME):
                    type_params["is_primary"] = True
                result_data["type_params"] = type_params

            if len(result.problems) > 0:
                result_data.yaml_set_comment_before_after_key(
                    key="name",
                    before=f"{ConfigFileRenderer.UNKNOWN_FIELD_VALUE}: " + ", ".join(result.problems),
                )

            rendered.append(result_data)

        return rendered

    def _get_comments_for_unknown_columns(self, results: List[InferenceResult]) -> List[str]:
        delim = "\n  - "
        return [
            self._fixme(result.column.column_name + delim + delim.join(result.problems))
            for result in results
            if result.type_node == InferenceSignalType.UNKNOWN
        ]

    def _get_configuration_data_for_table(self, table: SqlTable, results: List[InferenceResult]) -> Dict:
        data = CommentedMap(
            {
                "semantic_model": {
                    "name": table.table_name,
                    "node_relation": {
                        "alias": table.table_name,
                        "schema_name": table.schema_name,
                        "database": table.db_name,
                        "relation_name": table.sql,
                    },
                    "entities": self._render_entity_columns(results),
                    "dimensions": self._render_dimension_columns(results),
                    "measures": [],
                }
            }
        )

        header_comments = [self._fixme("Unreviewed inferred config file")] + self._get_comments_for_unknown_columns(
            results
        )
        data.yaml_set_comment_before_after_key(
            key="semantic_model",
            before="\n".join(header_comments),
        )

        return data

    def render(self, results: List[InferenceResult]) -> None:
        """Render the inference results to files in the configured directory."""
        os.makedirs(self.dir_path, exist_ok=True)

        results_by_table: Dict[SqlTable, List[InferenceResult]] = defaultdict(list)
        for result in results:
            results_by_table[result.column.table].append(result)

        for table, results in results_by_table.items():
            table_data = self._get_configuration_data_for_table(table, results)
            table_path = self._get_filename_for_table(table)
            with open(table_path, "w") as table_file:
                yaml.dump(table_data, table_file)
