from __future__ import annotations
from typing import Tuple

from relationalai.semantics.metamodel import builtins, ir, factory as f, types, visitor, helpers
from relationalai.semantics.metamodel.compiler import Pass, group_tasks
from relationalai.semantics.metamodel.util import OrderedSet
from relationalai.semantics.metamodel.util import FrozenOrderedSet
from relationalai.semantics.metamodel.typer.typer import is_primitive

class FormatOutputs(Pass):
    def __init__(self, use_rel: bool=False):
        super().__init__()
        self._use_rel = use_rel

    #--------------------------------------------------
    # Public API
    #--------------------------------------------------
    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
        wide_outputs = options.get("wide_outputs", False)
        return self.OutputRewriter(wide_outputs, self._use_rel).walk(model)

    class OutputRewriter(visitor.Rewriter):
        def __init__(self, wide_outputs: bool = False, use_rel: bool = False):
            super().__init__()
            self.wide_outputs = wide_outputs
            self._use_rel = use_rel

        def handle_logical(self, node: ir.Logical, parent: ir.Node):
            # Rewrite children first
            node = super().handle_logical(node, parent)

            groups = group_tasks(node.body, {
                "outputs": ir.Output,
            })

            # If no outputs, return as is
            if not groups["outputs"]:
                return node

            if self.wide_outputs:
                return adjust_wide_outputs(node, groups["outputs"])

            return adjust_gnf_outputs(self._use_rel, node, groups["outputs"])

#--------------------------------------------------
# GNF vs wide output support
#--------------------------------------------------

# For wide outputs, only adjust the output task to include the keys.
# output looks like: (key0, key1, val0, val1, ...)
def adjust_wide_outputs(task: ir.Logical, outputs: OrderedSet[ir.Task]):
    body = list(task.body)
    for output in outputs:
        assert(isinstance(output, ir.Output))
        if output.keys:
            body.remove(output)
            body.append(rewrite_wide_output(output))
    return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)

# For GNF outputs we need to generate a rule for each "column" in the output
# and potentially one wide key column
def adjust_gnf_outputs(use_rel: bool, task: ir.Logical, outputs: OrderedSet[ir.Task]):
    body = list(task.body)
    for output in outputs:
        assert(isinstance(output, ir.Output))
        if output.keys:
            # Remove the original output. This is replaced by per-column outputs below
            body.remove(output)

            is_export = helpers.is_export(output)

            # Exports and Rel execution rely on all columns being in GNF format.
            if is_export or use_rel:
                _adjust_all_gnf_outputs(body, output, is_export)
            else: # Otherwise, put all keys into one wide keys relation
                _adjust_outputs_with_wide_keys(body, output)

    return ir.Logical(task.engine, task.hoisted, tuple(body), task.annotations)

# Generate an output for each "column"
# output looks like: def output(:cols, :col000, key0, key1, value)
def _adjust_all_gnf_outputs(body, output: ir.Output, is_export: bool):
    assert output.keys

    original_cols = OrderedSet()

    # Track existing column names (case-insensitively) to avoid value-column/key-column collisions
    # in exports, where Snowflake treats column names as case-insensitive.  e.g. value col "Person"
    # and key col "person" would collide as Snowflake column names.
    existing_names: set[str] = set()

    for idx, alias in enumerate(output.aliases):
        # Skip None values which are used as a placeholder for missing values
        if alias[1] is None:
            continue
        original_cols.add(alias[1])
        if is_export:
            existing_names.add(alias[0].lower())
        body.extend(_generate_output_column_gnf(output, idx, alias, is_export))

    idx = len(output.aliases)
    for key in output.keys:
        if key not in original_cols:
            col_name = key.name
            
            # If the key column name collides with an existing column name (case-insensitively),
            # append a "_key" suffix, then "_key_1", "_key_2", etc. until unique.
            # e.g. key "v" colliding with value col "v"  → "v_key"
            #      key "v" colliding with "v" and "v_key" → "v_key_1"
            if is_export:
                if col_name.lower() in existing_names:
                    col_name = f"{key.name}_key"
                    suffix = 1
                    while col_name.lower() in existing_names:
                        col_name = f"{key.name}_key_{suffix}"
                        suffix += 1
                existing_names.add(col_name.lower())

            body.extend(_generate_output_column_gnf(output, idx, (col_name, key), is_export))
            idx += 1

# Generate an output for each value "column" and one wide output for all the keys
#  * value output looks like: def output(:cols, :col000, key0, key1, value)
#  * key output looks like:
#    def output(:keys, output_key_0, output_key_1, other_key_0, ...)
#
# Exceptions: keys for exports and compound keys are converted to GNF, same as the
# value columns.
def _adjust_outputs_with_wide_keys(body, output: ir.Output):
    assert output.keys

    original_cols = OrderedSet()
    val_cols: list[Tuple[str, ir.Value] | None] = []
    key_cols: OrderedSet[Tuple[str, ir.Value]] = OrderedSet()
    key_cols.add(("keys", f.literal("keys", types.Symbol))) # name key col so we can identify it later
    for alias in output.aliases:
        # None values are used as a placeholder for missing values
        # They are added to maintain the correct col count when enumerated below
        if alias[1] is None:
            val_cols.append(None)
            continue

        original_cols.add(alias[1])

        if isinstance(alias[1], ir.Var) and alias[1] in output.keys: # note: skips compound keys
            key_cols.add(alias)
        else:
            val_cols.append(alias)

    # Add keys not in output to the end
    for key in output.keys:
        if key not in original_cols:
            key_cols.add((key.name, key))

    # Generate GNF val cols
    for idx, alias in enumerate(val_cols):
        if alias:
            new_col = _generate_output_column(output, idx, alias, key_cols)
            body.extend(new_col)

    # Create a wide key column with all keys
    if len(key_cols) > 1:
        body.append(ir.Output(
            output.engine,
            key_cols.frozen(),
            output.keys,
            output.annotations
        ))

# Generate a relation representing a single col in GNF form
def _generate_output_column(output: ir.Output, idx: int, alias: tuple[str, ir.Value], key_cols):
    if not output.keys:
        return [output]

    aliases = [("cols", f.literal("cols", types.Symbol))]
    aliases.append(("col", f.literal(f"col{idx:03}", types.Symbol)))

    # Append all keys at the start
    keys = iter(key_cols)
    assert next(keys) == ("keys", f.literal("keys", types.Symbol)) # skip col name
    for key in keys:
        aliases.append((f"key_{key[0]}_{idx}", key[1]))

    aliases.append(alias) # append val

    return [
        ir.Output(
            output.engine,
            FrozenOrderedSet.from_iterable(aliases),
            output.keys,
            output.annotations
        )
    ]

# Generate a relation representing a single col in GNF form for export
def _generate_output_column_gnf(output: ir.Output, idx: int, alias: tuple[str, ir.Value], is_export: bool):
    if not output.keys:
        return [output]

    aliases = [("cols", f.literal("cols", types.Symbol))] if not is_export else []
    aliases.append(("col", f.literal(f"col{idx:03}", types.Symbol)))

    for k in output.keys:
        aliases.append((f"key_{k.name}_{idx}", k))

    if (is_export and
        isinstance(alias[1], ir.Var) and
        (not is_primitive(alias[1].type) or alias[1].type == types.Hash)):

        # Choose the base name for the uuid variable:
        #  - If it is a key column: use alias[1].name (original variable name) to keep the uuid name
        #    clean, unaffected by collision-avoidance renaming on alias[0].
        #    e.g. key "v" renamed to "v_key" → uuid stays "v_3_uuid", not "v_key_3_uuid"
        #  - If it is a Hash-typed value column: use alias[0] (user-provided column name) so the
        #    uuid export column remains meaningful to the user.
        #    e.g. value col "my_hash" → uuid is "my_hash_2_uuid"
        uuid_base = alias[1].name if alias[1] in output.keys else alias[0]
        uuid = f.var(f"{uuid_base}_{idx}_uuid", types.String)

        if not is_primitive(alias[1].type):
            # For non-primitive types, we keep the original alias
            aliases.append((alias[0], uuid))
        else:
            # For Hash types, we use the uuid name as alias
            aliases.append((uuid.name, uuid))

        return [
            ir.Lookup(None, builtins.uuid_to_string, (alias[1], uuid)),
            ir.Output(
                output.engine,
                FrozenOrderedSet.from_iterable(aliases),
                output.keys,
                output.annotations
            )
        ]
    else:
        aliases.append(alias)

        return [
            ir.Output(
                output.engine,
                FrozenOrderedSet.from_iterable(aliases),
                output.keys,
                output.annotations
            )
        ]

def rewrite_wide_output(output: ir.Output):
    assert(output.keys)

    # Only append keys that are not already in the output
    suffix_keys = []
    for key in output.keys:
        if all([val is not key for _, val in output.aliases]):
            suffix_keys.append(key)

    aliases: OrderedSet[Tuple[str, ir.Value]] = OrderedSet()

    # Add the remaining args, unless it is already a key
    for name, val in output.aliases:
        if not isinstance(val, ir.Var) or val not in suffix_keys:
            aliases.add((name, val))

    # Add the keys to the output
    for key in suffix_keys:
        aliases.add((key.name, key))

    # TODO - we are assuming that the Rel compiler will translate nullable lookups
    # properly, returning a `Missing` if necessary, like this:
    # (nested_192(_adult, _adult_name) or (not nested_192(_adult, _) and _adult_name = Missing)) and
    return ir.Output(
        output.engine,
        aliases.frozen(),
        output.keys,
        output.annotations
    )

    # TODO: in the rel compiler, see if we can do this outer join
    # 1. number of keys
    # 2. each relation
    # 3. each variable, starting with the keys
    # 4. tag output with @arrow

    # @arrow def output(_book, _book_title, _author_name):
    #   rel_primitive_outer_join(#1, book_title, author_name, _book, _book_title, _author_name)
    # def output(p, n, c):
    #     rel_primitive_outer_join(#1, name, coolness, p, n, c)
