from __future__ import annotations

from pandas import DataFrame
from typing import Any, Union, Tuple, Literal, TYPE_CHECKING

from relationalai.clients.config import Config
from relationalai.semantics.metamodel import Model, Task, ir
from relationalai.semantics.metamodel.visitor import collect_by_type
if TYPE_CHECKING:
    from relationalai.semantics.internal.internal import Model as InternalModel

from .util import NameCache

import rich

# global flag to suppress type errors from being printed
SUPPRESS_TYPE_ERRORS = False

class Executor():
    """ Interface for an object that can execute the program specified by a model. """
    def execute(self, model: Model, task:Task, format:Literal["pandas", "snowpark"]="pandas",
                export_to: Any = None, update: bool = False, meta: dict[str, Any] | None = None,
                source_info: dict[str, Any] | None = None) -> Union[DataFrame, Any]:
        raise NotImplementedError(f"execute: {self}")

    def _compute_cols(self, task: ir.Task, final_model: ir.Model|None) -> Tuple[list[str], list[str], list[int]]:
        cols = [] # all cols in output
        extra_cols = [] # all key cols not in output
        key_locs = [] # locations of output key cols in all output cols

        # we assume only queries have outputs
        original_outputs = collect_by_type(ir.Output, task) if task else None
        outputs = collect_by_type(ir.Output, final_model) if final_model else None

        # there are some outputs, and some have keys
        if original_outputs and outputs and not all(not out.keys for out in outputs):
            assert len(original_outputs) == 1
            original_output = original_outputs[0]
            original_cols = []
            original_cols_val = []

            keys = outputs[0].keys
            assert keys

            for out in outputs:
                assert out.keys is not None
                assert set(out.keys) == set(keys), "outputs with different key sets in the same query"

            for (idx, (alias, val)) in enumerate(original_output.aliases):
                if not alias:
                    continue
                original_cols.append(alias)
                original_cols_val.append(val)
                if isinstance(val, ir.Var) and val in keys:
                    key_locs.append(idx)

            extra_cols = []
            name_cache = NameCache(start_from_one=True)
            for key in keys:
                if isinstance(key, ir.Var) and key not in original_cols_val:
                    extra_cols.append(f"hidden_{name_cache.get_name(key.id, key.name)}")
            cols = original_cols + extra_cols
        elif outputs:
            cols = [alias for alias, _ in outputs[-1].aliases if alias]

        return cols, extra_cols, key_locs

    def _postprocess_df(self, config: Config, df: DataFrame, extra_cols: list[str]) -> DataFrame:
        if bool(config.get("compiler.debug_hidden_keys", False)):
            rich.print("[blue]DataFrame with extra hidden columns:[/blue]")
            rich.print(f"[black]{df}[/black]")
        if not df.empty:
            for col in extra_cols:
                if col in df.columns:
                    df = df.drop(col, axis=1)
        return df

    def export_to_csv(self, model: "InternalModel", query) -> str:
        ### Only implemented in the LQP executor for now.
        raise NotImplementedError(f"export_to_csv is not supported by {type(self).__name__}")

    def batch_export_to_csv(self, model: "InternalModel", queries: list) -> list[str]:
        ### Only implemented in the LQP executor for now.
        raise NotImplementedError(f"batch_export_to_csv is not supported by {type(self).__name__}")
