import json
from collections.abc import Sequence
from os import PathLike
from typing import Optional, Literal, NamedTuple, Self, TypeVar, AnyStr, ClassVar, Any

import pydantic
from pydantic import ConfigDict, Field, model_validator, NonNegativeInt, TypeAdapter
from pydantic.alias_generators import to_camel

AnyPath = AnyStr | PathLike

Plugin = Literal["module", "declaration", "symbol", "elaboration"]

ALL_PLUGINS = ("module", "declaration", "symbol", "elaboration")


def plugin_short_name(plugin: Plugin) -> str:
    match plugin:
        case "module":
            return "mod"
        case "declaration":
            return "decl"
        case "symbol":
            return "sym"
        case "elaboration":
            return "elab"
        case _:
            raise ValueError


class BaseModel(pydantic.BaseModel):
    """
    :meta private:
    """
    model_config = ConfigDict(alias_generator=to_camel, populate_by_name=True, strict=True)


M = TypeVar("M", bound="RootModel")


class RootModel(BaseModel):
    _plugin_name: ClassVar[Plugin]

    @classmethod
    def from_obj(cls: type[M], obj: Any) -> list[M]:
        return TypeAdapter(list[cls]).validate_python(obj)

    @classmethod
    def from_str(cls: type[M], data: str) -> list[M]:
        """Construct a list of models from a str"""
        return TypeAdapter(list[cls]).validate_json(data)

    @classmethod
    def from_json_file(cls: type[M], filename: AnyPath) -> list[M]:
        """Read a list of models from a json file"""
        # Use from_obj instead of from_str since the latter can cause stack overflow on large objects
        with open(filename) as fp:
            obj = json.load(fp)
            return cls.from_obj(obj)


LeanName = Sequence[str | NonNegativeInt]


def is_internal(name: LeanName) -> bool:
    """A heuristic to find if this name is generated by the Lean compiler"""
    if any(type(x) is not str for x in name) or \
            any(name[-1].startswith(p) for p in ["_", "eq_", "match_", "proof_"]):
        return True
    return False


def pp_name(name: LeanName):
    """Pretty print a Lean name."""
    return ".".join(str(x) for x in name)


def parse_name(s: str) -> LeanName:
    return s.split(".")


T = TypeVar("T")


def is_prefix_of(p: Sequence[T], s: Sequence[T]) -> bool:
    return len(p) <= len(s) and all(x == y for x, y in zip(p, s))


class ModuleInfo(BaseModel):
    """Module-level information"""
    _plugin_name = "module"

    imports: list[LeanName]
    """Modules directly imported by this module"""
    docstring: list[str]
    """List of all docstrings in this module, as marked by /-! ... -/"""


# Declaration
BinderInfo = Literal["default", "implicit", "strictImplicit", "instImplicit"]
Visibility = Literal["regular", "protected", "private"]
RecKind = Literal["default", "partial", "nonrec"]
DeclarationKind = Literal[
    "abbrev",
    "axiom",
    "classInductive",
    "definition",
    "example",
    "inductive",
    "instance",
    "opaque",
    "structure",
    "theorem",
    "proofWanted",
]


def is_disjoint_union(*xs) -> bool:
    return sum(x is not None for x in xs) == 1

def ensure_disjoint_union(o, fields: list[str]):
    if not is_disjoint_union(getattr(o, f) for f in fields):
        raise TypeError("exactly one of [" + ", ".join(fields) + "] is expected")


class StringRange(NamedTuple):
    """A byte range within a file"""
    start: NonNegativeInt
    stop: NonNegativeInt

    def as_slice(self) -> slice:
        return slice(self.start, self.stop)


class Param(BaseModel):
    """A parameter to a declaration, as defined in the source code"""
    ref: Optional[StringRange]
    """The entire syntax object, e.g., `{x : A}`"""
    id: Optional[StringRange]
    """The identifier part, e.g., `x` in {x : A}"""
    type: Optional[StringRange]
    """The type part, e.g., `A` in {x : A}"""
    binder_info: BinderInfo = Field(alias="bi")
    """The binder kind of this parameter, e.g., `implicit` for {x : A}"""


class Modifiers(BaseModel):
    """A modifier attached to a declaration"""
    visibility: Visibility
    """Visibility level of a declaration"""
    is_noncomputable: bool
    rec_kind: RecKind
    """Recursion level of a declaration"""
    is_unsafe: bool
    docstring: Optional[str] = Field(default=None, alias="docString")


class Syntax(BaseModel):
    """A simplified representation of a Syntax node"""
    original: bool
    """Whether this Syntax node is original or generated (e.g., by a macro)"""
    range: Optional[StringRange]


class PPSyntax(Syntax):
    """A Syntax node with pretty-printing info"""
    pp: Optional[str] = Field(default=None)


class OpenDecl(BaseModel):
    """An `open` directive in Lean"""

    class Simple(BaseModel):
        namespace: LeanName
        hiding: list[LeanName] = Field(alias="except")

    class Rename(BaseModel):
        name: LeanName
        as_: LeanName = Field(alias="as")

    simple: Optional[Simple] = Field(default=None)
    rename: Optional[Rename] = Field(default=None)

    @model_validator(mode="after")
    def ensure_disjoint_union(self) -> Self:
        ensure_disjoint_union(self, ["simple", "rename"])
        return self


class ScopeInfo(BaseModel):
    """Current scope info, i.e., variables, namespaces, etc.  Used for isolating declarations in a file"""
    var_decls: list[str]
    """`variable` directives"""
    include_vars: list[LeanName]
    """variables marked with `include`"""
    omit_vars: list[LeanName]
    """variables marked with `omit`"""
    curr_namespace: LeanName
    """The current namespace"""
    open_decl: list[OpenDecl]
    """`open` directives"""


class Declaration(RootModel):
    """Declarations in the source code"""
    _plugin_name = "declaration"

    kind: DeclarationKind
    """Kind of a declaration"""
    ref: PPSyntax
    """"""
    name: LeanName
    signature: PPSyntax
    """The signature part, e.g., `{α : Sort u} (a : α) : α` in `def id {α : Sort u} (a : α) : α`"""
    modifiers: Modifiers
    params: list[Param]
    type: Optional[Syntax]
    value: Optional[Syntax]
    scope_info: ScopeInfo


# Symbol
SymbolKind = Literal[
    "axiom",
    "definition",
    "theorem",
    "opaque",
    "quotient",
    "inductive",
    "constructor",
    "recursor",
]


class Symbol(RootModel):
    """
    A symbol, as seen by the Lean kernel.

    Usually called a `constant` in Lean terminology.
    """
    _plugin_name = "symbol"

    kind: SymbolKind
    name: LeanName
    type: str
    type_references: list[LeanName]
    """Names used in defining the type of this symbol"""
    value_references: Optional[list[LeanName]]
    """Names used in defining the value of this symbol, or None if it has no value"""
    is_prop: bool
    """Whether the type of this symbol is a `Prop`"""


# Context / Goal
class Variable(BaseModel):
    """
    A variable in Lean contexts.

    Usually called a :term:`free variable` or an `fvar` in Lean terminology.
    """
    name: LeanName
    binder_info: Optional[BinderInfo] = Field(default=None)
    """Binder info of this variable, or None if it is defined by a `let`"""
    type: str
    value: Optional[str] = Field(default=None)
    """Value of this variable if it is defined by a `let`, or None otherwise"""
    is_prop: bool


Context = list[Variable]
"""A :term:`local context` in Lean"""


class Goal(BaseModel):
    """A :term:`metavariable` in Lean"""
    tag: LeanName
    """The user name of this metavariable"""
    context: Context
    type: str
    is_prop: bool


ProofState = list[Goal]
"""A :term:`proof state`, represented by a list of goals"""


# Elaboration
class PPSyntaxWithKind(PPSyntax):
    """A Syntax node with pretty-printing and kind info"""
    kind: LeanName
    """Syntax kind"""


class TacticElabInfo(BaseModel):
    """An InfoTree node about a tactic"""
    references: list[LeanName]
    """Names directly referenced in the tactic"""
    before: ProofState
    after: ProofState


class SpecialValue(BaseModel):
    """Marks the current value is of a special form of interest"""
    const: Optional[LeanName] = Field(default=None)
    """This value is a constant reference"""
    fvar: Optional[LeanName] = Field(default=None)
    """This value is a fvar reference"""

    @model_validator(mode="after")
    def ensure_disjoint_union(self) -> Self:
        ensure_disjoint_union(self, ["const", "fvar"])
        return self



class TermElabInfo(BaseModel):
    """An InfoTree node about a term"""
    context: Context
    type: str
    expected_type: Optional[str]
    value: str
    special: Optional[SpecialValue] = Field(default=None)


class MacroInfo(BaseModel):
    """An InfoTree node about a macro expansion"""
    expanded: PPSyntaxWithKind


SimpleElabInfo = Literal["command", "field", "option", "completion", "uw", "custom", "alias", "redecl", "omission", "partial", "term"]


class ElabInfo(BaseModel):
    """
    An InfoTree node

    Can be one of term node, tactic node, macro node, or simple node.

    A simple node is one with only its kind recorded.
    """
    term: Optional[TermElabInfo] = Field(default=None)
    tactic: Optional[TacticElabInfo] = Field(default=None)
    macro: Optional[MacroInfo] = Field(default=None)
    simple: Optional[SimpleElabInfo] = Field(default=None)

    @model_validator(mode="after")
    def ensure_disjoint_union(self) -> Self:
        ensure_disjoint_union(self, ["term", "tactic", "macro", "simple"])
        return self


class InfoTree(RootModel):
    _plugin_name = "elaboration"
    info: ElabInfo
    ref: PPSyntaxWithKind
    children: list[Self]
