from typing import Union
from dataclasses import dataclass, field
from contextlib import contextmanager
from contextvars import ContextVar
from relationalai.semantics.frontend.base import Statement, Relationship, Model, Concept, Chain
from relationalai.util.source import SourcePos

#------------------------------------------------------
# Monoid Types (only Boolean monoid supported at the moment)
#------------------------------------------------------

class Monoid():
    pass

class OrMonoid(Monoid):
    pass

@dataclass
class SumMonoid(Monoid):
    type: Concept

    def __post_init__(self):
        from relationalai.util.error import warn
        warn("Unstable monoid", "SumMonoid is currently unstable in the backend.")

@dataclass
class MaxMonoid(Monoid):
    type: Concept

@dataclass
class MinMonoid(Monoid):
    type: Concept

#------------------------------------------------------
# Instruction Types
#------------------------------------------------------

@dataclass
class Instruction:
    source: SourcePos = field(default_factory=SourcePos.new, kw_only=True)

@dataclass
class Assign(Instruction):
    expr: Statement

@dataclass
class Empty(Instruction):
    relationship: Relationship | Chain

@dataclass
class InstructionWithArity(Instruction):
    arity: int

@dataclass
class Upsert(InstructionWithArity):
    expr: Statement

@dataclass
class InstructionWithMonoid(InstructionWithArity):
    monoid: Monoid

@dataclass
class MonoidInstr(InstructionWithMonoid):
    expr: Statement

@dataclass
class MonusInstr(InstructionWithMonoid):
    expr: Statement

@dataclass
class Break(Instruction):
    conditions: tuple[Statement, ...]
    log_every: int | None = None

#------------------------------------------------------
# Context Variables for tracking current builder
#------------------------------------------------------

# The current instruction container (ProcedureBuilder or LoopBuilder)
_current_container: ContextVar[Union["ProcedureBuilder", "LoopBuilder", None]] = ContextVar(
    'current_container', default=None
)

def _register_instruction(instr: Instruction) -> None:
    """Register an instruction with the current container if in context."""
    container = _current_container.get()
    if container is None:
        raise RuntimeError("Instructions can only be used inside an procedure")
    elif isinstance(container, ProcedureBuilder) or isinstance(container, LoopBuilder):
        container._instructions.append(instr)

#------------------------------------------------------
# InstructionContext helper class
#------------------------------------------------------

class InstructionContext:
    """Represents the context behind instructions that behave like insert/deletes."""

    def __init__(self, arity: int | None = None, monoid=None):
        self.arity = arity
        self._monoid = monoid

    def with_arity(self, arity: int):
        assert self.arity is None, "Cannot redefine arity of an instruction"
        return InstructionContext(arity=arity, monoid=self._monoid)

    def with_monoid(self, monoid: Monoid):
        assert self._monoid is None, "Cannot redefine monoid of an instruction"
        return InstructionContext(arity=self.arity, monoid=monoid)

    def upsert(self, expr: Statement) -> Upsert:
        assert self.arity is not None, "Upsert must have an integer arity. Use `with_arity(x).upsert(...)`"
        instr = Upsert(arity=self.arity, expr=expr)
        _register_instruction(instr)
        return instr

    def monoid(self, expr: Statement) -> MonoidInstr:
        assert self.arity is not None, "Monoid must have an integer arity. Use `with_arity(x).monoid(...)`"
        assert self._monoid is not None, "Monoid must have a monoid. Use `with_monoid(x).monoid(...)`"
        instr = MonoidInstr(arity=self.arity, monoid=self._monoid, expr=expr)
        _register_instruction(instr)
        return instr

    def monus(self, expr: Statement) -> MonusInstr:
        assert self.arity is not None, "Monus must have an integer arity. Use `with_arity(x).monus(...)`"
        assert self._monoid is not None, "Monus must have a monoid. Use `with_monoid(x).monus(...)`"
        instr = MonusInstr(arity=self.arity, monoid=self._monoid, expr=expr)
        _register_instruction(instr)
        return instr

#------------------------------------------------------
# LoopBuilder
#------------------------------------------------------

class LoopBuilder(Instruction):
    """Represents a loop in the engine."""

    def __init__(self, over: list[Union[Concept, Relationship]] = []):
        # TODO: Currently looping over things is not supported.
        self.source = SourcePos.new()
        self._over = over
        self._instructions: list[Instruction] = []

    def __str__(self):
        return self._format(indent=0)

    def _format(self, indent: int = 0) -> str:
        prefix = "  " * indent
        lines = [f"{prefix}Loop:"]
        for instr in self._instructions:
            lines.append(self._format_instruction(instr, indent + 1))
        return "\n".join(lines)

    def _format_instruction(self, instr: Instruction, indent: int) -> str:
        prefix = "  " * indent
        if isinstance(instr, LoopBuilder):
            return instr._format(indent)
        elif isinstance(instr, Assign):
            return f"{prefix}@assign {instr.expr}"
        elif isinstance(instr, Empty):
            return f"{prefix}@empty {instr.relationship}"
        elif isinstance(instr, MonoidInstr):
            return f"{prefix}@monoid:{instr.arity}:{instr.monoid} {instr.expr}"
        elif isinstance(instr, MonusInstr):
            return f"{prefix}@monus:{instr.arity}:{instr.monoid} {instr.expr}"
        elif isinstance(instr, Upsert):
            return f"{prefix}@upsert:{instr.arity} {instr.expr}"
        elif isinstance(instr, Break):
            le = f"(log_every={instr.log_every}) " if instr.log_every is not None else ""
            return f"{prefix}@break{le} {instr.conditions}"
        else:
            return f"{prefix}{instr}"

#------------------------------------------------------
# ProcedureBuilder
#------------------------------------------------------

class ProcedureBuilder:
    """Builds an imperative procedure with global relationships and instructions.

    An ProcedureBuilder represents a complete procedure containing:
    - Global relationships (loop heads)
    - A sequence of instructions (assign, empty, upsert, break, or loops)
    """

    def __init__(self, model: Model, *globals: Relationship | Chain):
        self._model = model
        self._globals: tuple[Relationship | Chain, ...] = globals
        self._instructions: list[Instruction] = []

    def assign(self, expr: Statement) -> Assign:
        instr = Assign(expr=expr)
        _register_instruction(instr)
        return instr

    def empty(self, relationship: Relationship | Chain) -> Empty:
        instr = Empty(relationship=relationship)
        _register_instruction(instr)
        return instr

    def break_if(self, *conditions: Statement, log_every: int | None = None) -> Break:
        instr = Break(conditions=conditions, log_every=log_every)
        _register_instruction(instr)
        return instr

    def with_arity(self, arity: int) -> InstructionContext:
        return InstructionContext(arity=arity)

    def with_monoid(self, monoid: Monoid) -> InstructionContext:
        return InstructionContext(monoid=monoid)

    def loop(self, over: list[Union[Concept, Relationship]] = []):
        return loop(over)

    def __str__(self):
        lines = ["Procedure:"]

        # Print globals
        if self._globals:
            lines.append("  Globals:")
            for rel in self._globals:
                lines.append(f"    {rel}")

        # Print instructions
        if self._instructions:
            lines.append("  Instructions:")
            for instr in self._instructions:
                lines.append(self._format_instruction(instr, indent=2))

        return "\n".join(lines)

    def _format_instruction(self, instr: Instruction, indent: int) -> str:
        prefix = "  " * indent
        if isinstance(instr, LoopBuilder):
            return instr._format(indent)
        elif isinstance(instr, Assign):
            return f"{prefix}@assign {instr.expr}"
        elif isinstance(instr, Empty):
            return f"{prefix}@empty {instr.relationship}"
        elif isinstance(instr, MonoidInstr):
            return f"{prefix}@monoid:{instr.arity}:{instr.monoid} {instr.expr}"
        elif isinstance(instr, MonusInstr):
            return f"{prefix}@monus:{instr.arity}:{instr.monoid} {instr.expr}"
        elif isinstance(instr, Upsert):
            return f"{prefix}@upsert:{instr.arity} {instr.expr}"
        elif isinstance(instr, Break):
            le = f"(log_every={instr.log_every}) " if instr.log_every is not None else ""
            return f"{prefix}@break{le} {instr.conditions}"
        else:
            return f"{prefix}{instr}"

#------------------------------------------------------
# Context Managers
#------------------------------------------------------

@contextmanager
def procedure(*globals: Relationship | Chain, model: Model | None = None):
    """Context manager for defining an procedure with global relationships.

    Usage:
        with procedure(Reachable, Frontier):
            assign(...)
            empty(...)
            with loop():
                upsert(...)
                break_if(...)

    Parameters
    ----------
    *globals : Relationship
        Relationships that are modified by the procedure. Important to specify for compilation.

    Yields
    ------
    ProcedureBuilder
        The procedure builder (can be used to access the built structure).
    """
    from ... import _check_model
    model = _check_model() if model is None else model
    builder = ProcedureBuilder(model, *globals)

    model._procedure_fragment(builder)

    token = _current_container.set(builder)
    try:
        yield builder
    finally:
        _current_container.reset(token)


@contextmanager
def loop(over: list[Union[Concept, Relationship]] = []):
    """Context manager for defining a loop within an procedure.

    Usage:
        with procedure(...):
            with loop():
                # instructions inside the loop (while-loop)
                assign(...)
                break_if(...)

        with procedure(...):
            with loop(over=SomeRelation):
                # instructions inside the loop (for-loop over relation)
                ...

    Can be nested for inner loops:
        with procedure(...):
            with loop():
                with loop():
                    # nested loop
                    ...

    Parameters
    ----------
    over : optional
        Relationship to iterate over (for-loop). If None, creates a while-loop.
    """
    parent = _current_container.get()
    if parent is None:
        raise RuntimeError("loop() must be used inside procedure() or another loop()")

    loop_builder = LoopBuilder(over=over)
    # Register this loop as an instruction in the parent container
    parent._instructions.append(loop_builder)

    token = _current_container.set(loop_builder)
    try:
        yield loop_builder
    finally:
        _current_container.reset(token)
