from dataclasses import dataclass, field
from typing import Callable, Union, TYPE_CHECKING
import juic

if TYPE_CHECKING:
    import juic.eval

# Mapping of Jui types to Python types:
#   Null           None
#   Booleans       bool
#   Integers       int
#   Floats         float
#   Strings        str
#   C/C++ refs     CXXQualid
#   List           list
#   "this"         ThisRef
#   Functions      Function, BuiltinFunction
#   Record types   RecordType
#   Records        Record

@dataclass
class CXXQualid:
    qualid: str

@dataclass
class BuiltinFunction:
    func: Callable

@dataclass
class Function:
    # Context in which the function gets evaluated
    closure: "juic.eval.Closure"
    # Expression node to evaluate when calling the function
    body: "juic.parser.Node"
    # Parameter names, must all be unique. May be empty
    params: list[str] = field(default_factory=list)
    # Name of variadic argument if one; must also be unique
    variadic: str | None = None

@dataclass
class RecordType:
    # Record name
    name: str
    # Base type for inheritance/subconstructor logic
    base: Union[None, "RecordType"]

    # TODO: Record prototypes?

@dataclass
class RecordCtor:
    func: Function

@dataclass
class Record:
    # A record type if it's pure, or a thunk with the base object otherwise.
    base: Union[RecordType, "juic.eval.Thunk"]
    # Standard key-value attribute pairs
    attr: dict[str, Union["JuiValue", "juic.eval.Thunk"]]
    # Children elements
    children: list[Union["JuiValue", "juic.eval.Thunk"]]
    # TODO: Keep track of variables that are not fields, i.e. "methods"?
    # scope: dict[str, "JuiValue"]
    # TODO: Labels
    # labels: dict[str, "JuiValue"]

JuiValue = Union[None, bool, int, float, str, CXXQualid, list,
                 "juic.eval.PartialRecordSnapshot", BuiltinFunction, Function,
                 RecordType, RecordCtor, Record]

def juiIsArith(value):
    return type(value) in [bool, int, float]

def juiIsAddable(value):
    return juiIsArith(value) or type(value) in [str, list]

def juiIsComparable(value):
    return type(value) in [None, bool, int, float, str]

def juiIsLogical(value):
    return type(value) == bool

def juiIsUnpackable(value):
    return type(value) == list

def juiIsCallable(value):
    return type(value) in [BuiltinFunction, Function]

def juiIsProjectable(value):
    return type(value) in [Record, juic.eval.PartialRecordSnapshot]

def juiIsConstructible(value):
    return type(value) in [RecordType, RecordCtor]

def juiIsRecordUpdatable(value):
    return type(value) == Record

# Generic functions on values

# String representation of a value
def juiValueString(v):
    match v:
        case None:
            return "null"
        case bool():
            return str(v).lower()
        case int() | float():
            return str(v)
        case str():
            return repr(v)
        case CXXQualid():
            return "&" + v.qualid
        case list():
            return "[" + ", ".join(juiValueString(x) for x in v) + "]"
        case BuiltinFunction():
            return str(v)
        case Function() as f:
            p = f.params + (["..." + f.variadic] if f.variadic else [])
            s = "fun(" + ", ".join(p) + ") => " + str(f.body)
            s += " (in some closure)"
            return s
        case RecordCtor() as rc:
            f = rc.func
            p = f.params + (["..." + f.variadic] if f.variadic else [])
            s = "rec(" + ", ".join(p) + ") => " + str(f.body)
            s += " (in some closure)"
            return s
        case RecordType() as rt:
            return str(rt)
        case Record() as r:
            s = r.base.name + " {"
            s += "; ".join(x + ": " + juiValueString(y) for x, y in r.attr.items())
            if len(r.attr) and len(r.children):
                s += "; "
            s += "; ".join(juiValueString(x) for x in r.children)
            return s + "}"
            raise NotImplementedError
        case juic.eval.Thunk() as th:
            return str(th)
        case juic.eval.PartialRecordSnapshot() as prs:
            return str(prs)
        case _:
            raise NotImplementedError

# Check whether two *forced* values are equal
def juiValuesEqual(v1, v2):
    def unwrapThunk(v):
        match v:
            case juic.eval.Thunk() as th:
                assert th.evaluated and not th.invalid
                return unwrapThunk(th.result)
            case _:
                return v

    v1 = unwrapThunk(v1)
    v2 = unwrapThunk(v2)
    match v1, v2:
        case None, None:
            return True
        case bool(), bool():
            return v1 == v2
        case int(), int():
            return v1 == v2
        case float(), float():
            return v1 == v2
        case str(), str():
            return v1 == v2
        case CXXQualid(), CXXQualid():
            return v1.qualid == v2.qualid
        case list() as l1, list() as l2:
            return len(l1) == len(l2) and \
                all(juiValuesEqual(v1, v2) for v1, v2 in zip(l1, l2))
        case BuiltinFunction(), BuiltinFunction():
            return id(v1) == id(v2)
        case Function(), Function():
            raise Exception("cannot compare functions")
        case RecordType(), RecordType():
            return id(v1) == id(v2)
        case Record() as r1, Record() as r2:
            if not juiValuesEqual(r1.base, r2.base):
                return False
            if r1.attr.keys() != r2.attr.keys():
                return False
            if any(not juiValuesEqual(r1.attr[k], r2.attr[k]) for k in r1.attr):
                return False
            if len(r1.children) != len(r2.children):
                return False
            if any(not juiValuesEqual(v1, v2)
                   for v1, v2 in zip(r1.children, r2.children)):
                return False
            return True
        case juic.eval.PartialRecordSnapshot(), \
             juic.eval.PartialRecordSnapshot():
            raise Exception("cannot compare PRSs")
        case _, _:
            raise Exception(f"invalid types for comparison: {type(v1)}, {type(v2)}")