from __future__ import annotations
import trace
from typing import TYPE_CHECKING, Any, Dict

if TYPE_CHECKING:  # to prevent circular import
    from opto.trace.nodes import Node
from opto.trace.bundle import bundle
import copy


@bundle()
def clone(x: Any):
    """This is a clone operator of x."""
    return copy.deepcopy(x)



[docs]
def identity(x: Any):
    # identity(x) behaves the same as x.clone()
    return x.clone()



# Unary operators and functions


@bundle()
def pos(x: Any):
    """This is a pos operator of x."""
    return +x


@bundle()
def neg(x: Any):
    """This is a neg operator of x."""
    return -x


@bundle()
def abs(x: Any):
    """This is an abs operator of x."""
    return abs(x)


@bundle()
def invert(x: Any):
    """This is an invert operator of x."""
    return ~x


@bundle()
def round(x: Any, n: Any):
    """This is a round operator of x."""
    return round(x, n)


@bundle()
def floor(x: Any):
    """This is a floor operator of x."""
    import math

    return math.floor(x)


@bundle()
def ceil(x: Any):
    """This is a ceil operator of x."""
    import math

    return math.ceil(x)


@bundle()
def trunc(x: Any):
    """This is a trunc operator of x."""
    import math

    return math.trunc(x)


# Normal arithmetic operators


@bundle()
def add(x: Any, y: Any):
    """This is an add operator of x and y."""
    return x + y


@bundle()
def subtract(x: Any, y: Any):
    """This is a subtract operator of x and y."""
    return x - y


@bundle()
def multiply(x: Any, y: Any):
    """This is a multiply operator of x and y."""
    return x * y


@bundle()
def floor_divide(x: Any, y: Any):
    """This is a floor_divide operator of x and y."""
    return x // y


@bundle()
def divide(x: Any, y: Any):
    """This is a divide operator of x and y."""
    return x / y


@bundle()
def mod(x: Any, y: Any):
    """This is a mod operator of x and y."""
    return x % y


@bundle()
def node_divmod(x: Any, y: Any):
    """This is a divmod operator of x and y."""
    return divmod(x, y)


@bundle()
def power(x: Any, y: Any):
    """This is a power operator of x and y."""
    return x**y


@bundle()
def lshift(x: Any, y: Any):
    """This is a lshift operator of x and y."""
    return x << y


@bundle()
def rshift(x: Any, y: Any):
    """This is a rshift operator of x and y."""
    return x >> y


@bundle()
def and_(x: Any, y: Any):
    """This is an and operator of x and y."""
    return x & y


@bundle()
def or_(x: Any, y: Any):
    """This is an or operator of x and y."""
    return x | y


@bundle()
def xor(x: Any, y: Any):
    """This is a xor operator of x and y."""
    return x ^ y


# Comparison methods


@bundle()
def lt(x: Any, y: Any):
    """This is a lt operator of x and y."""
    return x < y


@bundle()
def le(x: Any, y: Any):
    """This is a le operator of x and y."""
    return x <= y


@bundle()
def eq(x: Any, y: Any):
    """This is an eq operator of x and y."""
    return x == y


@bundle()
def neq(x: Any, y: Any):
    """This is a not eq operator of x and y."""
    return x != y


@bundle()
def ne(x: Any, y: Any):
    """This is a ne operator of x and y."""
    return x != y


@bundle()
def ge(x: Any, y: Any):
    """This is a ge operator of x and y."""
    return x >= y


@bundle()
def gt(x: Any, y: Any):
    """This is a gt operator of x and y."""
    return x > y


# logical operators


@bundle()
def cond(condition: Any, x: Any, y: Any):
    """This selects x if condition is True, otherwise y."""
    x, y, condition = x, y, condition  # This makes sure all data are read
    return x if condition else y


@bundle()
def not_(x: Any):
    """This is a not operator of x."""
    return not x


@bundle()
def is_(x: Any, y: Any):
    """Whether x is equal to y."""
    return x is y


@bundle()
def is_not(x: Any, y: Any):
    """Whether x is not equal to y."""
    return x is not y


@bundle()
def in_(x: Any, y: Any):
    """Whether x is in y."""
    return x in y


@bundle()
def not_in(x: Any, y: Any):
    """Whether x is not in y."""
    return x not in y


# Indexing and slicing
@bundle()
def getitem(x: Any, index: Any):
    """This is a getitem operator of x based on index."""
    return x[index]


@bundle()
def pop(x: Any, index: Any):
    """This is a pop operator of x based on index."""
    return x.pop(index)


@bundle()
def len_(x: Any):
    """This is a len operator of x."""
    return len(x)


# String operators
@bundle()
def ord_(x: Any):
    """The unicode number of a character."""
    return ord(x)


@bundle()
def chr_(x: Any):
    """The character of a unicode number."""
    return chr(x)


@bundle()
def concat(x: Any, y: Any):
    """This is a concatenation operator of x and y."""
    return x + y


@bundle()
def lower(x: Any):
    """This makes all characters in x lower case."""
    return x.lower()


@bundle()
def upper(x: Any):
    """This makes all characters in x upper case."""
    return x.upper()


@bundle()
def title(x: Any):
    """This makes the first character to upper case and the rest to lower case."""
    return x.title()


@bundle()
def swapcase(x: Any):
    """Swaps the case of all characters: uppercase character to lowercase and vice-versa."""
    return x.swapcase()


@bundle()
def capitalize(x: Any):
    """Converts the first character of a string to uppercase."""
    return x.capitalize()


@bundle()
def split(x: Any, y: Any, maxsplit: Any = -1):
    """Splits the string by finding a substring y in string x, return the first part and second part of string x without y."""
    return x.split(y, maxsplit)


@bundle()
def strip(x: Any, chars=None):
    """Removes the leading and trailing characters of x."""
    return x.strip(chars)


@bundle()
def replace(x: Any, old: Any, new: Any, count: Any = -1):
    """Replaces all occurrences of substring y in string x with z."""
    return x.replace(old, new, count)


@bundle()
def format(x: Any, *args, **kwargs):
    """Fills in a string template with content, str.format()."""
    return x.format(*args, **kwargs)


@bundle()
def join(x: Any, *y: Any):
    """Joins a sequence y with different strs with x: "\n".join(["a", "b", "c"]) -> "a\nb\nc"."""
    return x.join(y)


@bundle()
def node_getattr(obj: Node, attr: str):
    """This operator gets attr of obj."""
    return getattr(obj, attr)


@bundle(
    _process_inputs=False,
    allow_external_dependencies=True,
)
def call(fun: Node, *args, **kwargs):
    """This operator calls the function `fun` with args (args_0, args_1, etc.) and kwargs. If there are no args or kwargs, i.e. call(fun=function_name), the function takes no input."""
    # Run the function as it is
    fun = fun._data
    # Call the node with the input arguments
    assert callable(fun), "The function must be callable."
    output = fun(*args, **kwargs)
    return output


@bundle()
def to_list(x: Any):
    """This converts x to a list."""
    return list(x)


@bundle()
def make_list(*args):
    """This creates a list from the arguments."""
    return list(args)


@bundle()
def to_dict(x: Any):
    """This converts x to a dictionary."""
    return dict(x)


@bundle()
def make_dict(**kwargs):
    """This creates a dictionary from the keyword arguments."""
    return kwargs


@bundle()
def to_set(x: Any):
    """This converts x to a set."""
    return set(x)


@bundle()
def make_set(*args):
    """This creates a set from the arguments."""
    return set(args)


@bundle()
def to_tuple(x: Any):
    """This converts x to a tuple."""
    return tuple(x)


@bundle()
def make_tuple(*args):
    """This creates a tuple from the arguments."""
    return tuple(args)


# dict operators


@bundle()
def keys(x: Dict):
    """Return the keys of a dictionary x as a list."""
    if not isinstance(x, dict):
        raise AttributeError(f"{type(x)} object has no attribute 'values'.")

    return [k for k in x.keys()]


@bundle()
def values(x: Dict):
    """Return the values of a dictionary x as a list."""
    if not isinstance(x, dict):
        raise AttributeError(f"{type(x)} object has no attribute 'values'.")

    return [k for k in x.values()]


# dict in-place operators


@bundle()
def dict_update(x: Dict, y: Dict):
    """Update the dictionary x with the dictionary y."""
    x = copy.copy(x)
    x.update(y)
    return x


@bundle()
def dict_pop(x: Dict, key: Any):
    """Pop the key from the dictionary x."""
    x = copy.copy(x)
    x.pop(key)
    return x


@bundle()
def dict_popitem(x: Dict):
    """Pop the last item from the dictionary x."""
    x = copy.copy(x)
    x.popitem()
    return x


# list in-place operators


@bundle()
def list_append(x: Any, y: Any):
    """Append y to x."""
    x = copy.copy(x)
    x.append(y)
    return x


@bundle()
def list_clear(x: Any):
    """Clear x."""
    x = copy.copy(x)
    x.clear()
    return x


@bundle()
def list_extend(x: Any, y: Any):
    """Extend x with y."""
    x = copy.copy(x)
    x.extend(y)
    return x


@bundle()
def list_insert(x: Any, index: Any, y: Any):
    """Insert y at index in x."""
    x = copy.copy(x)
    x.insert(index, y)
    return x


@bundle()
def list_pop(x: Any, index: Any):
    """Pop the index from x."""
    x = copy.copy(x)
    x.pop(index)
    return x


@bundle()
def list_remove(x: Any, y: Any):
    """Remove y from x."""
    x = copy.copy(x)
    x.remove(y)
    return x


@bundle()
def list_reverse(x: Any):
    """Reverse x."""
    x = copy.copy(x)
    x.reverse()
    return x


@bundle()
def list_sort(x: Any, key: Any = None, reverse: Any = False):
    """Sort x."""
    x = copy.copy(x)
    x.sort(key=key, reverse=reverse)
    return x


# set in-place operators
@bundle()
def set_add(x: Any, y: Any):
    """Add y to x."""
    x = copy.copy(x)
    x.add(y)
    return x


@bundle()
def set_clear(x: Any):
    """Clear x."""
    x = copy.copy(x)
    x.clear()
    return x


@bundle()
def set_discard(x: Any, y: Any):
    """Discard y from x."""
    x = copy.copy(x)
    x.discard(y)
    return x


@bundle()
def set_intersection_update(x: Any, y: Any):
    """Update x with the intersection of x and y."""
    x = copy.copy(x)
    x.intersection_update(y)
    return x


@bundle()
def set_pop(x: Any):
    """Pop an element from x."""
    x = copy.copy(x)
    x.pop()
    return x


@bundle()
def set_remove(x: Any, y: Any):
    """Remove y from x."""
    x = copy.copy(x)
    x.remove(y)
    return x


@bundle()
def set_symmetric_difference_update(x: Any, y: Any):
    """Update x with the symmetric difference of x and y."""
    x = copy.copy(x)
    x.symmetric_difference_update(y)
    return x


@bundle()
def set_update(x: Any, y: Any):
    """Update x with y."""
    x = copy.copy(x)
    x.update(y)
    return x


@bundle()
def call_llm(system_prompt, *user_prompts, **kwargs):
    """Query the language model of system_prompt with user_prompts."""
    messages = [{"role": "system", "content": system_prompt}]
    for user_prompt in user_prompts:
        messages.append({"role": "user", "content": user_prompt})
    from opto.utils.llm import AutoGenLLM

    llm = AutoGenLLM()
    response = llm(messages=messages, **kwargs)
    return response.choices[0].message.content