from functools import reduce
from itertools import chain
from dataclasses import dataclass, field
from typing import OrderedDict, cast as _c

from pyparsing import ParseResults

from . import (
    Name,
    Param,
    Decl,
    ir,
    grammar as _g,
    fresh,
    Def,
    Example,
    Ctor,
    Data,
    Sig,
    Field,
    Class,
    Instance,
)


@dataclass(frozen=True)
class Node:
    loc: int


@dataclass(frozen=True)
class Type(Node): ...


@dataclass(frozen=True)
class Ref(Node):
    name: Name


@dataclass(frozen=True)
class FnType(Node):
    param: Param[Node]
    ret: Node


@dataclass(frozen=True)
class Fn(Node):
    param: Name
    body: Node


@dataclass(frozen=True)
class Call(Node):
    callee: Node
    arg: Node
    implicit: str | bool


@dataclass(frozen=True)
class Placeholder(Node):
    is_user: bool


@dataclass(frozen=True)
class Nomatch(Node):
    arg: Node


@dataclass(frozen=True)
class Case(Node):
    ctor: Ref
    params: list[Name]
    body: Node


@dataclass(frozen=True)
class Match(Node):
    arg: Node
    cases: list[Case]


_g.name.add_parse_action(lambda r: Name(r[0][0]))

_ops = {"+": "add", "-": "sub", "*": "mul", "/": "div"}


def _infix(loc: int, ret: ParseResults):
    r = ret[0]
    if not isinstance(r, ParseResults):
        return r
    return Call(loc, Call(loc, Ref(loc, Name(_ops[r[1]])), r[0], False), r[2], False)


_g.expr.add_parse_action(_infix)
_g.type_.add_parse_action(lambda l, r: Type(l))
_g.ph.add_parse_action(lambda l, r: Placeholder(l, True))
_g.ref.add_parse_action(lambda l, r: Ref(l, r[0][0]))
_g.i_param.add_parse_action(lambda r: Param(r[0], r[1], True))
_g.e_param.add_parse_action(lambda r: Param(r[0], r[1], False))
_g.c_param.add_parse_action(lambda r: Param(r[0], r[1], True, True))
_g.fn_type.add_parse_action(lambda l, r: FnType(l, r[0], r[1]))
_g.fn.add_parse_action(lambda l, r: reduce(lambda a, n: Fn(l, n, a), reversed(r[0]), r[1]))
_g.match.add_parse_action(lambda l, r: Match(l, r[0], list(r[1])))
_g.case.add_parse_action(lambda r: Case(r[0].loc, r[0], r[1], r[2]))
_g.nomatch.add_parse_action(lambda l, r: Nomatch(l, r[0][0]))
_g.i_arg.add_parse_action(lambda l, r: (r[1], r[0]))
_g.e_arg.add_parse_action(lambda l, r: (r[0], False))
_g.call.add_parse_action(lambda l, r: reduce(lambda a, b: Call(l, a, b[0], b[1]), r[1:], r[0]))
_g.p_expr.add_parse_action(lambda r: r[0])

_g.return_type.add_parse_action(lambda l, r: r[0] if len(r) else Placeholder(l, False))
_g.def_.add_parse_action(lambda r: Def(r[0].loc, r[0].name, list(r[1]), r[2], r[3]))
_g.example.add_parse_action(lambda l, r: Example(l, list(r[0]), r[1], r[2]))
_g.type_arg.add_parse_action(lambda r: (r[0], r[1]))
_g.ctor.add_parse_action(lambda r: Ctor(r[0].loc, r[0].name, list(r[1]), list(r[2])))
_g.data.add_condition(lambda r: r[0].name.text == r[3], message="open and datatype name mismatch").add_parse_action(
    lambda r: Data(r[0].loc, r[0].name, list(r[1]), list(r[2]))
)
_g.c_field.add_parse_action(lambda l, r: Field(l, r[0], r[1]))
_g.class_.add_condition(lambda r: r[0].name.text == r[3], message="open and class name mismatch").add_parse_action(
    lambda r: Class(r[0].loc, r[0].name, list(r[1]), list(r[2]))
)
_g.i_field.add_parse_action(lambda r: (r[0], r[1]))
_g.inst.add_parse_action(lambda l, r: Instance(l, r[0], list(r[1])))


@dataclass(frozen=True)
class Parser:
    is_markdown: bool = False

    def __ror__(self, s: str):
        if not self.is_markdown:
            return list(_g.program.parse_string(s, parse_all=True))
        return chain.from_iterable(r[0] for r in _g.markdown.scan_string(s))


class DuplicateVariableError(Exception): ...


class UndefinedVariableError(Exception): ...


@dataclass(frozen=True)
class NameResolver:
    locals: dict[str, Name] = field(default_factory=dict)
    globals: dict[str, Name] = field(default_factory=dict)

    def __ror__(self, decls: list[Decl]):
        return [self._decl(d) for d in decls]

    def _decl(self, decl: Decl) -> Decl:
        self.locals.clear()

        if isinstance(decl, Def) or isinstance(decl, Data):
            self._insert_global(decl.loc, decl.name)

        if isinstance(decl, Def) or isinstance(decl, Example):
            return self._def_or_example(decl)

        if isinstance(decl, Data):
            return self._data(decl)

        if isinstance(decl, Class):
            return self._class(decl)

        return self._inst(_c(Instance, decl))

    def _def_or_example(self, d: Def[Node] | Example[Node]):
        params = self._params(d.params)
        ret = self.expr(d.ret)
        body = self.expr(d.body)
        if isinstance(d, Example):
            return Example(d.loc, params, ret, body)
        return Def(d.loc, d.name, params, ret, body)

    def _data(self, d: Data[Node]):
        params = self._params(d.params)
        ctors = [self._ctor(c, d.name) for c in d.ctors]
        return Data(d.loc, d.name, params, ctors)

    def _ctor(self, c: Ctor[Node], ty_name: Name):
        params = self._params(c.params)
        ty_args = [(self.expr(n), self.expr(t)) for n, t in c.ty_args]
        for p in params:
            del self.locals[p.name.text]
        self._insert_global(c.loc, c.name)
        return Ctor(c.loc, c.name, params, ty_args, ty_name)

    def _class(self, c: Class[Node]):
        params = self._params(c.params)
        fields = []
        for f in c.fields:
            self._insert_global(f.loc, f.name)
            fields.append(Field(f.loc, f.name, self.expr(f.type)))
        self._insert_global(c.loc, c.name)
        return Class(c.loc, c.name, params, fields)

    def _inst(self, i: Instance[Node]):
        t = self.expr(i.type)
        fields = []
        field_ids = set()
        for n, v in i.fields:
            n = _c(Ref, self.expr(n))
            if n.name.id in field_ids:
                raise DuplicateVariableError(n.name.text, n.loc)
            field_ids.add(n.name.id)
            fields.append((n, (self.expr(v))))
        return Instance(i.loc, t, fields)

    def _params(self, params: list[Param[Node]]):
        ret = []
        for p in params:
            self._insert_local(p.name)
            ret.append(Param(p.name, self.expr(p.type), p.is_implicit, p.is_class))
        return ret

    def expr(self, n: Node) -> Node:
        if isinstance(n, Ref):
            if v := self.locals.get(n.name.text, self.globals.get(n.name.text)):
                return Ref(n.loc, v)
            raise UndefinedVariableError(n.name.text, n.loc)
        if isinstance(n, FnType):
            typ = self.expr(n.param.type)
            b = self._with_locals(n.ret, n.param.name)
            p = Param(n.param.name, typ, n.param.is_implicit, n.param.is_class)
            return FnType(n.loc, p, b)
        if isinstance(n, Fn):
            return Fn(n.loc, n.param, self._with_locals(n.body, n.param))
        if isinstance(n, Call):
            return Call(n.loc, self.expr(n.callee), self.expr(n.arg), n.implicit)
        if isinstance(n, Nomatch):
            return Nomatch(n.loc, self.expr(n.arg))
        if isinstance(n, Match):
            arg = self.expr(n.arg)
            cases = []
            for c in n.cases:
                ctor = _c(Ref, self.expr(c.ctor))
                body = self._with_locals(c.body, *c.params)
                cases.append(Case(c.loc, ctor, c.params, body))
            return Match(n.loc, arg, cases)
        assert isinstance(n, Type) or isinstance(n, Placeholder)
        return n

    def _with_locals(self, node: Node, *names: Name):
        olds = [(v, self._insert_local(v)) for v in names]
        ret = self.expr(node)
        for v, old in olds:
            if old:
                self._insert_local(old)
            elif not v.is_unbound():
                del self.locals[v.text]
        return ret

    def _insert_local(self, v: Name):
        if v.is_unbound():
            return None
        old = self.locals.get(v.text)
        self.locals[v.text] = v
        return old

    def _insert_global(self, loc: int, name: Name):
        if not name.is_unbound():
            if name.text in self.globals:
                raise DuplicateVariableError(name.text, loc)
            self.globals[name.text] = name


class TypeMismatchError(Exception): ...


class UnsolvedPlaceholderError(Exception): ...


class UnknownCaseError(Exception): ...


class DuplicateCaseError(Exception): ...


class CaseParamMismatchError(Exception): ...


class CaseMissError(Exception): ...


class FieldMissError(Exception): ...


class UnknownFieldError(Exception): ...


@dataclass(frozen=True)
class TypeChecker:
    globals: dict[int, Decl] = field(default_factory=dict)
    locals: dict[int, Param[ir.IR]] = field(default_factory=dict)
    holes: OrderedDict[int, ir.Hole] = field(default_factory=OrderedDict)
    recur_ids: set[int] = field(default_factory=set)

    def __ror__(self, ds: list[Decl]):
        ret = [self._run(d) for d in ds]
        for i, h in self.holes.items():
            if h.answer.is_unsolved():
                ty = self._inliner().run(h.answer.type)
                if _is_solved_class(ty):
                    continue
                p = ir.Placeholder(i, h.is_user)
                raise UnsolvedPlaceholderError(str(p), h.locals, ty, h.loc)
        return ret

    def _run(self, decl: Decl) -> Decl:
        self.locals.clear()
        if isinstance(decl, Def) or isinstance(decl, Example):
            return self._def_or_example(decl)
        if isinstance(decl, Data):
            return self._data(decl)
        if isinstance(decl, Instance):
            return self._inst(decl)
        return self._class(_c(Class, decl))

    def _def_or_example(self, d: Def[Node] | Example[Node]):
        params = self._params(d.params)
        ret = self.check(d.ret, ir.Type())

        if isinstance(d, Def):
            self.globals[d.name.id] = Sig(d.loc, d.name, params, ret)
        body = self.check(d.body, ret)

        if isinstance(d, Example):
            return Example(d.loc, params, ret, body)

        checked = Def(d.loc, d.name, params, ret, body)
        self.globals[d.name.id] = checked
        return checked

    def _data(self, d: Data[Node]):
        params = self._params(d.params)
        data = Data(d.loc, d.name, params, [])
        self.globals[d.name.id] = data
        data.ctors.extend(self._ctor(c) for c in d.ctors)
        return data

    def _ctor(self, c: Ctor[Node]):
        params = self._params(c.params)
        ty_args: list[tuple[ir.IR, ir.IR]] = []
        for x, v in c.ty_args:
            x_val, x_ty = self.infer(x)
            v_val = self.check(v, x_ty)
            ty_args.append((x_val, v_val))
        ctor = Ctor(c.loc, c.name, params, ty_args, c.ty_name)
        self.globals[c.name.id] = ctor
        return ctor

    def _class(self, c: Class[Node]):
        params = self._params(c.params)
        fs = [Field(f.loc, f.name, self.check(f.type, ir.Type()), c.name) for f in c.fields]
        self.globals.update({f.name.id: f for f in fs})
        cls = Class(c.loc, c.name, params, fs)
        self.globals[c.name.id] = cls
        return cls

    def _inst(self, i: Instance[Node]):
        ty = self.check(i.type, ir.Type())
        if not isinstance(ty, ir.Class):
            raise TypeMismatchError("class", str(ty), i.type.loc)
        c = _c(Class, self.globals[ty.name.id])
        vals = {_c(Ref, n).name.id: (n, f) for n, f in i.fields}
        fields = []
        for f in c.fields:
            nv = vals.pop(f.name.id, None)
            if not nv:
                raise FieldMissError(f.name.text, i.loc)
            f_decl = _c(Field, self.globals[f.name.id])
            field_ty = ir.from_field(f_decl, c, False)[1]
            env = []
            for ty_arg in ty.args:
                assert isinstance(field_ty, ir.FnType)
                env.append((field_ty.param.name, ty_arg))
                field_ty = field_ty.ret
            f_type = self._inliner().run_with(field_ty, *env)
            fields.append((ir.Ref(f.name), self.check(nv[1], f_type)))
        for n, _ in vals.values():
            assert isinstance(n, Ref)
            raise UnknownFieldError(c.name.text, n.name.text, n.loc)
        c.instances.append(i.id)
        inst = Instance(i.loc, _c(ir.IR, ty), fields, i.id)
        self.globals[i.id] = inst
        return inst

    def _params(self, params: list[Param[Node]]):
        ret = []
        for p in params:
            t = self.check(p.type, ir.Type())
            if p.is_class:
                t = self._inliner().run(t)
                assert p.is_implicit
                if not isinstance(t, ir.Class):
                    raise TypeMismatchError("class", str(t), p.type.loc)
            param = Param(p.name, t, p.is_implicit, p.is_class)
            self.locals[p.name.id] = param
            ret.append(param)
        return ret

    def check(self, n: Node, typ: ir.IR) -> ir.IR:
        if isinstance(n, Fn):
            t = self._inliner().run(typ)
            if not isinstance(t, ir.FnType):
                raise TypeMismatchError(str(t), "function", n.loc)
            ret = self._inliner().run_with(t.ret, (t.param.name, ir.Ref(n.param)))
            p = Param(n.param, t.param.type, t.param.is_implicit, t.param.is_class)
            return ir.Fn(p, self._check_with(n.body, ret, p))

        val, got = self.infer(n)
        got = self._inliner().run(got)
        want = self._inliner().run(typ)

        if _can_insert_placeholders(want):
            holes_len = len(self.holes)
            if new_f := _with_placeholders(n, got, False):
                [self.holes.popitem() for _ in range(len(self.holes) - holes_len)]
                val, got = self.infer(new_f)

        if not self._eq(got, want):
            raise TypeMismatchError(str(want), str(got), n.loc)

        return val

    def infer(self, n: Node) -> tuple[ir.IR, ir.IR]:
        if isinstance(n, Ref):
            if param := self.locals.get(n.name.id):
                return ir.Ref(param.name), param.type
            d = self.globals[n.name.id]
            if isinstance(d, Def):
                return ir.from_def(d)
            if isinstance(d, Sig):
                self.recur_ids.add(d.name.id)
                return ir.from_sig(d)
            if isinstance(d, Data):
                return ir.from_data(d)
            if isinstance(d, Ctor):
                _, v, ty = ir.from_ctor(d, _c(Data, self.globals[d.ty_name.id]))
                return v, ty
            if isinstance(d, Field):
                return ir.from_field(d, _c(Class, self.globals[d.cls_name.id]))
            return ir.from_class(_c(Class, d))
        if isinstance(n, FnType):
            p_typ = self.check(n.param.type, ir.Type())
            p = Param(n.param.name, p_typ, n.param.is_implicit, n.param.is_class)
            b_val = self._check_with(n.ret, ir.Type(), p)
            return ir.FnType(p, b_val), ir.Type()
        if isinstance(n, Call):
            holes_len = len(self.holes)
            f_val, got = self.infer(n.callee)

            if implicit_f := _with_placeholders(n.callee, got, n.implicit):
                [self.holes.popitem() for _ in range(len(self.holes) - holes_len)]
                return self.infer(Call(n.loc, implicit_f, n.arg, n.implicit))

            if not isinstance(got, ir.FnType):
                raise TypeMismatchError("function", str(got), n.callee.loc)

            x_tm = self._check_with(n.arg, got.param.type, got.param)
            typ = self._inliner().run_with(got.ret, (got.param.name, x_tm))
            val = self._inliner().apply(f_val, x_tm)
            return val, typ
        if isinstance(n, Placeholder):
            ty = self._insert_hole(n.loc, n.is_user, ir.Type())
            v = self._insert_hole(n.loc, n.is_user, ty)
            return v, ty
        if isinstance(n, Nomatch):
            _, got = self.infer(n.arg)
            if not isinstance(got, ir.Data):
                raise TypeMismatchError("datatype", str(got), n.arg.loc)
            data = _c(Data, self.globals[got.name.id])
            for c in data.ctors:
                self._exhaust(n.arg.loc, c, data, got)
            return ir.Nomatch(), self._insert_hole(n.loc, False, ir.Type())
        if isinstance(n, Match):
            return self._infer_match(n)
        assert isinstance(n, Type)
        return ir.Type(), ir.Type()

    def _infer_match(self, n: Match):
        arg, arg_ty = self.infer(n.arg)
        if not isinstance(arg_ty, ir.Data):
            raise TypeMismatchError("datatype", str(arg_ty), n.arg.loc)
        data = _c(Data, self.globals[arg_ty.name.id])
        ctors = {c.name.id: c for c in data.ctors}
        ty: ir.IR | None = None
        cases: dict[int, ir.Case] = {}
        for c in n.cases:
            ctor = ctors.get(c.ctor.name.id)
            if not ctor:
                raise UnknownCaseError(data.name.text, c.ctor.name.text, c.loc)
            if ctor.name.id in cases:
                raise DuplicateCaseError(ctor.name.text, c.loc)
            c_params, c_ty = self._case_params(c.loc, ctor, data)
            if not self._eq(c_ty, arg_ty):
                raise TypeMismatchError(str(arg_ty), str(c_ty), c.loc)
            if len(c.params) != len(c_params):
                raise CaseParamMismatchError(len(ctor.params), len(c.params), c.loc)
            ps = [Param(n, p.type, False) for n, p in zip(c.params, c_params)]
            if ty is None:
                body, ty = self._infer_with(c.body, *ps)
            else:
                body = self._check_with(c.body, ty, *ps)
            cases[ctor.name.id] = ir.Case(ctor.name, ps, body)
        for c in [c for i, c in ctors.items() if i not in cases]:
            self._exhaust(n.loc, c, data, arg_ty)
        return ir.Match(arg, cases), ty

    def _inliner(self):
        return ir.Inliner(self.holes, self.globals, self.locals)

    def _eq(self, got: ir.IR, want: ir.IR):
        return ir.Converter(self.holes, self.globals, self.locals).eq(got, want)

    def _check_with(self, n: Node, typ: ir.IR, *ps: Param[ir.IR]):
        self.locals.update({p.name.id: p for p in ps})
        ret = self.check(n, typ)
        [self.locals.pop(p.name.id, None) for p in ps]
        return ret

    def _infer_with(self, n: Node, *ps: Param[ir.IR]):
        self.locals.update({p.name.id: p for p in ps})
        v, ty = self.infer(n)
        [self.locals.pop(p.name.id, None) for p in ps]
        return v, ty

    def _insert_hole(self, loc: int, is_user: bool, typ: ir.IR):
        i = fresh()
        self.holes[i] = ir.Hole(loc, is_user, self.locals.copy(), ir.Answer(typ))
        return ir.Placeholder(i, is_user)

    def _case_params(self, loc: int, c: Ctor[ir.IR], d: Data[ir.IR]):
        miss, v, ty = ir.from_ctor(c, d)
        while isinstance(ty, ir.FnType):
            p = ty.param
            x = self._insert_hole(loc, False, p.type) if p.is_implicit else ir.Ref(p.name)
            ty = self._inliner().run_with(ty.ret, (p.name, x))

            if miss > 0:
                miss -= 1
                assert isinstance(v, ir.Fn)
                v = self._inliner().run_with(v.body, (v.param.name, x))

        params = []
        while isinstance(v, ir.Fn):
            params.append(v.param)
            v = v.body

        return params, ty

    def _exhaust(self, loc: int, c: Ctor[ir.IR], d: Data[ir.IR], want: ir.IR):
        with ir.dirty_holes(self.holes):
            if self._eq(self._case_params(loc, c, d)[1], want):
                raise CaseMissError(c.name.text, loc)


def _is_solved_class(ty: ir.IR):
    if isinstance(ty, ir.Class):
        return all(not isinstance(a, ir.Ref) for a in ty.args)
    return False


def _can_insert_placeholders(ty: ir.IR):
    return not isinstance(ty, ir.FnType) or not ty.param.is_implicit


def _with_placeholders(f: Node, f_ty: ir.IR, implicit: str | bool) -> Node | None:
    if not isinstance(f_ty, ir.FnType):
        return None

    if isinstance(implicit, bool):
        if not f_ty.param.is_implicit:
            return None
        return _call_placeholder(f) if not implicit else None

    pending = 0
    while True:
        if not isinstance(f_ty, ir.FnType) or not f_ty.param.is_implicit:
            raise UndefinedVariableError(implicit, f.loc)
        if f_ty.param.name.text == implicit:
            break
        pending += 1
        f_ty = f_ty.ret

    if not pending:
        return None

    for _ in range(pending):
        f = _call_placeholder(f)
    return f


def _call_placeholder(f: Node):
    return Call(f.loc, f, Placeholder(f.loc, False), True)


check_string = lambda s, md=False: s | Parser(md) | NameResolver() | TypeChecker()
