from __future__ import annotations

from typing import Optional, Sequence as PySequence

from v0.relationalai.semantics.metamodel import ir, compiler as c, factory as f, util
from v0.relationalai.semantics.metamodel.helpers import collect_implicit_vars
from v0.relationalai.semantics.lqp.algorithms import is_script, mk_assign, is_instruction

class Splinter(c.Pass):
    """
    Splits multi-headed rules into multiple rules.
    """

    def __init__(self):
        super().__init__()
        self.name_cache = util.NameCache(start_from_one=True)

    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
        if isinstance(model.root, ir.Logical):
            final = []
            new_relations:list[ir.Relation] = []
            new_relations.extend(model.relations)
            for child in model.root.body:
                result, relations = self.handle(child)
                if isinstance(result, list):
                    final.extend(result)
                else:
                    final.append(result)
                new_relations.extend(relations)
            return ir.Model(
                model.engines,
                util.FrozenOrderedSet.from_iterable(new_relations),
                model.types,
                ir.Logical(
                    model.root.engine,
                    model.root.hoisted,
                    tuple(final)
                )
            )
        return model

    def handle(self, task: ir.Task) -> tuple[ir.Task | list[ir.Task], list[ir.Relation]]:
        if isinstance(task, ir.Logical):
            return self.handle_logical(task)
        elif isinstance(task, ir.Sequence):
            return self.handle_sequence(task)
        elif isinstance(task, ir.Loop):
            return self.handle_loop(task)
        return task, []

    def handle_logical(self, task: ir.Logical) -> tuple[ir.Task | list[ir.Task], list[ir.Relation]]:
        # If this Logical wraps a Sequence, recurse into children without splitting the wrapper
        has_sequence = any(isinstance(c, ir.Sequence) for c in task.body)
        if has_sequence:
            new_body: list[ir.Task] = []
            new_relations: list[ir.Relation] = []
            for child in task.body:
                result, relations = self.handle(child)
                if isinstance(result, list):
                    new_body.extend(result)
                else:
                    new_body.append(result)
                new_relations.extend(relations)
            return ir.Logical(task.engine, task.hoisted, tuple(new_body), task.annotations), new_relations
        # No Sequence inside — split as usual
        split_result, relation = self.split(task)
        new_relations = [relation] if relation is not None else []
        return (split_result if len(split_result) > 1 else split_result[0]), new_relations

    def handle_sequence(self, task: ir.Sequence) -> tuple[ir.Sequence, list[ir.Relation]]:
        new_tasks: list[ir.Task] = []
        new_relations: list[ir.Relation] = []
        for child in task.tasks:
            result, relations = self.handle(child)
            if isinstance(result, list):
                for i in result:
                    if isinstance(i, ir.Logical) and not is_instruction(i):
                        assigned = mk_assign(i)
                        assert isinstance(assigned, ir.Task)
                        new_tasks.append(assigned)
                    else:
                        new_tasks.append(i)
            else:
                if isinstance(result, ir.Logical) and not is_instruction(result):
                    assigned = mk_assign(result)
                    assert isinstance(assigned, ir.Task)
                    new_tasks.append(assigned)
                else:
                    new_tasks.append(result)
            new_relations.extend(relations)
        return ir.Sequence(task.engine, task.hoisted, tuple(new_tasks), task.annotations), new_relations

    def handle_loop(self, task: ir.Loop) -> tuple[ir.Loop, list[ir.Relation]]:
        result, relations = self.handle(task.body)
        assert isinstance(result, ir.Task) and is_script(result)
        body = result
        if body is task.body:
            return task, relations
        return ir.Loop(task.engine, task.hoisted, task.iter, body, task.concurrency, task.annotations), relations

    def split(self, node: ir.Logical) -> tuple[list[ir.Task], Optional[ir.Relation]]:
        # Split this logical, which represents a rule, into potentially many logicals, one
        # for each head (update or output)
        effects, body = self.split_items(node.body)

        if len(effects) > 1:
            effects_vars = collect_implicit_vars(*effects)
            final:list[ir.Task] = []
            connection = None
            if body:
                # if the node has a body, create a connection, derive the body into the
                # connection, and then lookup this connection on the other logicals
                name = self.name_cache.get_name(node.id, "_intermediate")
                connection = f.relation(name, [f.field(f"f_{i+1}", v.type) for i, v in enumerate(effects_vars)])
                final.append(f.logical([*body, f.derive(connection, list(effects_vars))]))

            for effect in effects:
                if connection:
                    # there's a connection, create a logical that looks up from the connection
                    # and derives into the effect
                    lookup = f.lookup(connection, effects_vars.get_list())
                    final.append(f.logical([lookup, effect]))
                else:
                    # if there's no connection, the effect is self contained (probably literals)
                    final.append(f.logical([effect]))
            return final, connection
        return [node], None


    def split_items(self, items: PySequence[ir.Task]) -> tuple[list[ir.Task], list[ir.Task]]:
        effects = []
        body = []
        for item in items:
            if isinstance(item, (ir.Update, ir.Output)):
                effects.append(item)
            else:
                body.append(item)
        return effects, body
