from . import Pass
from ....metamodel import metamodel as mm
from ....metamodel.builtins import builtins as b
from ....metamodel.rewriter import Rewriter
from ....metamodel.metamodel_analyzer import VarFinder
from .....config import Config

from typing import List, cast


class TransformRequire(Pass):
    """
    Transforms `Require` nodes from the old format (unique_fields in check) to the new
    format (relation lookup in domain, unique in check). Also adds error tasks to each
    transformed `Require` node.

    Old format:
        Require
            domain
                Logical
            check
                Logical
                    unique_fields((person.ssn))
                    unique_fields((person.id_high, person.id_low))

    New format:
        Require
            domain
                Logical
                    Person_ssn(person::Person, ssn::String)
            check
                Logical
                    unique((ssn))
            error
                Logical
                    construct(Error, "message"::String, "Requirements not met"::String, "person"::String, person::Person, "ssn"::String, ssn::String, "pyrel_id"::String, 1185::Number(38,0), error::Error)
                    → derive Error(error::Error)
                    → derive message(error::Error, "Requirements not met"::String)
                    → derive person(error::Error, person::Person)
                    → derive ssn(error::Error, ssn::String)
                    → derive pyrel_id(error::Error, 1185::Number(38,0))
        Require
            domain
                Logical
                    Person_id(person::Person, id_low::Int, id_high::Int)
            check
                Logical
                    unique((id_low, id_high))
            error
                Logical
                    construct(Error, "message"::String, "Requirements not met"::String, "person"::String, person::Person, "id_low"::String, id_low::Int, "id_high"::String, id_high::Int, "pyrel_id"::String, 1185::Number(38,0), error::Error)
                    → derive Error(error::Error)
                    → derive message(error::Error, "Requirements not met"::String)
                    → derive person(error::Error, person::Person)
                    → derive id_low(error::Error, id_low::Int)
                    → derive id_high(error::Error, id_high::Int)
                    → derive pyrel_id(error::Error, 1185::Number(38,0))
    """

    def _rewrite_task(self, node: mm.Task, config: Config) -> mm.Task:
        transformer = self.RequireTransformer()
        return transformer(node)

    class RequireTransformer(Rewriter):
        def logical(self, node: mm.Logical):
            """
            Transform any Require nodes in the Logical body.
            """
            new_body = []

            for task in node.body:
                if isinstance(task, mm.Require):
                    new_requires = transform_require_node(task)
                    for req in new_requires:
                        # Ensure error is populated
                        req_with_error = req.mut(error=_generate_error(req))
                        # TODO: currently we wrap the result in a `Logical`, because the v0
                        # `Flatten` pass expects it to be nested in this way.
                        new_body.append(mm.Logical(body=(req_with_error,)))
                else:
                    new_body.append(task)

            return node.mut(body=tuple(new_body))


def _generate_error(node: mm.Require) -> mm.Task:
    """Generate synthetic error for a Require node.

    Creates a V1 error structure containing:
    - Construct: Creates Error entity with unique ID
    - Update nodes deriving:
      - Error entity (into CoreEntities.Error relation)
      - "message" attribute
      - Domain variables as attributes
      - "pyrel_id" attribute (using node.id for uniqueness)
    """

    if node.error is not None:
        return node.error

    # Extract domain variables
    domain_vars = VarFinder().find_vars(node.domain)

    # Create error variable and use node.id as pyrel_id
    error_type = cast(mm.ScalarType, b.CoreEntities.Error)
    error_var = mm.Var(type=error_type, name="error")
    pyrel_id = node.id  # Use node's unique ID

    # Helper to create attribute relations
    def create_error_attr_relation(attr_name: str, attr_type: mm.Type) -> mm.Relation:
        return mm.Relation(
            name=f"{attr_name}",
            fields=(
                mm.Field(name="error", type=error_type),  # type: ignore
                mm.Field(name=attr_name, type=attr_type)
            )
        )

    # Build Construct values to build the actual error entity
    values: list[mm.Value] = [
        error_type,
        mm.Literal(b.core.String, "message"),
        mm.Literal(b.core.String, "Requirements not met"),
    ]

    for var in domain_vars:
        values.extend([
            mm.Literal(b.core.String, var.name),
            var
        ])

    values.extend([
        mm.Literal(b.core.String, "pyrel_id"),
        mm.Literal(b.core.Integer, pyrel_id)
    ])

    construct = mm.Construct(values=tuple(values), id_var=error_var)

    # Create Update nodes which derive each attribute for the error
    updates = [
        mm.Update(
            relation=error_type,
            args=(error_var,),
            effect=mm.Effect.derive
        ),
        mm.Update(
            relation=create_error_attr_relation("message", b.core.String),
            args=(error_var, mm.Literal(b.core.String, "Requirements not met")),
            effect=mm.Effect.derive
        )
    ]

    for var in domain_vars:
        updates.append(mm.Update(
            relation=create_error_attr_relation(var.name, var.type),
            args=(error_var, var),
            effect=mm.Effect.derive
        ))

    updates.append(mm.Update(
        relation=create_error_attr_relation("pyrel_id", b.core.Integer),
        args=(error_var, mm.Literal(b.core.Integer, pyrel_id)),
        effect=mm.Effect.derive
    ))

    # Assemble error task
    error_task = mm.Logical(
        body=(construct, *updates),
        reasoner=node.reasoner,
        source=node.source
    )

    return error_task

def transform_require_node(node: mm.Require) -> List[mm.Require]:
    """
    Transform a single Require node from old format to new format.
    Returns a list of Require nodes.
    """
    # Only transform if check is a Logical node
    if not isinstance(node.check, mm.Logical):
        return [node]

    # Only transform if domain is empty Logical
    if not isinstance(node.domain, mm.Logical) or node.domain.body:
        return [node]

    # Look for unique_fields lookups in the check
    new_requires = []
    leftover_checks = []

    for item in node.check.body:
        if isinstance(item, mm.Lookup) and item.relation is b.constraints.unique_fields:
            # Extract fields from unique_fields((f1, f2, ...))
            fields = item.args[0]
            if not isinstance(fields, tuple) or not all(isinstance(f, mm.Field) for f in fields):
                leftover_checks.append(item)
                continue

            # Type narrowing: we know all elements are mm.Field now
            first_field = fields[0]
            assert isinstance(first_field, mm.Field)

            # Skip union types for now
            if isinstance(first_field.type, mm.UnionType):
                leftover_checks.append(item)
                continue

            # Get the relation from the first field
            relation = first_field._relation

            # Create variables for all fields in the relation
            all_vars = []
            unique_vars = []

            for field in relation.fields:
                var = mm.Var(type=field.type, name=field.name + "_var")
                all_vars.append(var)
                if field in fields:
                    unique_vars.append(var)

            # Create the new format:
            # - domain: R(v1, v2, v3, ...)
            # - check: unique((v1, v2)) where v1, v2 are the key fields
            domain = mm.Logical(body=(mm.Lookup(relation, tuple(all_vars)),))
            check_inner = mm.Lookup(b.constraints.unique, (tuple(unique_vars),))
            check = mm.Logical((check_inner,), source=node.check.source)

            new_req = mm.Require(
                domain=domain,
                check=check,
                error=node.error,
                annotations=node.annotations,
                reasoner=node.reasoner,
                source=node.source
            )
            new_requires.append(new_req)
        else:
            leftover_checks.append(item)

    # If we transformed some unique_fields, return the transformed requires
    if new_requires:
        # If there are leftover checks, create a Require for them
        if leftover_checks:
            leftover_req = mm.Require(
                check=mm.Logical(tuple(leftover_checks), source=node.check.source),
                annotations=node.annotations,
                reasoner=node.reasoner,
                source=node.source
            )
            new_requires.append(leftover_req)

        return new_requires

    return [node]
