"""
Nexa Language — Static Type Checker v0.2 (upgraded from v25)

Checks:
  - Undefined variable / function usage
  - Function call argument count mismatches
  - Type annotation mismatches (basic)
  - Null safety: assigning null to non-nullable annotated variable
  - Generics: <T> type param tracking — v0.2: type variables propagate into function body scope
  - let immutability: errors on reassignment of let-bound names (aligned with interpreter)
  - Trait/impl conformance: class implements all required trait methods
  - VarStatement and NativeAssertStatement awareness
  - GenericTypeAnnotation nodes: type params registered in scope
  - GenericRuntimeCheck: warns that runtime-only generics lack compile-time safety
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
from dataclasses import dataclass, field

from nexa_lang.ast_nodes import (
    Program, FunctionDef, ClassDef, LetStatement, Assignment,
    FunctionCall, ReturnStatement, IfExpr, ForStatement,
    WhileStatement, ThrowStatement, Identifier, BinaryOp,
    InterfaceDef, RangeExpr, MemberAccess,
    VarStatement, NativeAssertStatement, TraitDef, TraitImpl, UseStatement,
)

# ── Diagnostic ────────────────────────────────────────────────────────────────
@dataclass
class TypeDiag:
    level: str     # "error" | "warning"
    line: int
    message: str
    hint: str = ""

    def format_with_source(self, lines: List[str]) -> str:
        icon = "❌" if self.level == "error" else "⚠️ "
        res = f"  {icon}  [{self.line}]  {self.message}"
        if self.line > 0 and self.line <= len(lines):
            src_line = lines[self.line - 1]
            res += f"\n    | {self.line} | {src_line}"
            res += f"\n    | {' ' * len(str(self.line))} | ^~~~"
        if self.hint:
            res += f"\n    = help: {self.hint}"
        return res

    def __str__(self):
        icon = "❌" if self.level == "error" else "⚠️ "
        res = f"  {icon}  [{self.line}]  {self.message}"
        if self.hint:
            res += f"\n    = help: {self.hint}"
        return res

# ── Scope tracking ────────────────────────────────────────────────────────────
class TypeScope:
    def __init__(self, parent=None):
        self._vars: Dict[str, Optional[str]] = {}   # name -> type hint
        self._funcs: Dict[str, int] = {}             # name -> param count
        self._let_bindings: set = set()              # v25: names bound with let (immutable intent)
        self._traits: Dict[str, dict] = {}           # v25: trait name -> {methods: [...]}
        self._type_vars: Dict[str, str] = {}         # v0.2: generic type params e.g. T -> "type_param"
        self.parent = parent

    def declare(self, name: str, type_hint: Optional[str] = None, is_let: bool = False):
        self._vars[name] = type_hint
        if is_let:
            self._let_bindings.add(name)

    def declare_type_var(self, name: str):
        """v0.2: Register a generic type parameter (T, U, E, etc.) in this scope."""
        self._type_vars[name] = "TypeParam"
        self._vars[name] = "TypeParam"  # also visible as identifier

    def is_type_var(self, name: str) -> bool:
        if name in self._type_vars: return True
        return self.parent.is_type_var(name) if self.parent else False

    def is_let_bound(self, name: str) -> bool:
        """Return True if name was declared with let (immutable intent)."""
        if name in self._let_bindings:
            return True
        return self.parent.is_let_bound(name) if self.parent else False

    def declare_func(self, name: str, arity: int):
        self._funcs[name] = arity

    def lookup(self, name: str) -> bool:
        if name in self._vars or name in self._funcs: return True
        return self.parent.lookup(name) if self.parent else False

    def func_arity(self, name: str) -> Optional[int]:
        if name in self._funcs: return self._funcs[name]
        return self.parent.func_arity(name) if self.parent else None

    def type_of(self, name: str) -> Optional[str]:
        if name in self._vars: return self._vars[name]
        return self.parent.type_of(name) if self.parent else None


# ── Type Checker ──────────────────────────────────────────────────────────────
class TypeChecker:
    def __init__(self):
        self.diags: List[TypeDiag] = []
        self._interfaces: Dict[str, List[str]] = {}

        # Pre-populate built-in names so they don't raise "undefined"
        self._global = TypeScope()
        for name in [
            "print", "input", "range", "len", "type", "int", "float",
            "str", "bool", "list", "set", "zip", "sum", "max", "min",
            "abs", "round", "sorted", "reversed", "enumerate",
            "keys", "values", "items", "split", "join", "upper", "lower",
            "strip", "replace", "contains", "startswith", "endswith",
            "format", "isinstance", "append", "pop", "map", "filter",
            "reduce", "staticmethod", "true", "false", "null",
            "PI", "E", "INF",
            # All builtin modules
            "automl", "nlp", "vision", "math", "stats", "fs", "json",
            "http", "time", "random", "agent", "string", "array",
            "df", "db", "re", "serve", "tensor", "env", "sys",
            "datetime", "crypto", "base64", "task",
            # Common assertion helpers
            "assert_eq", "assert_neq", "assert_true", "assert_false",
            "assert_gt", "assert_lt", "assert_gte", "assert_lte",
            "assert_nil", "assert_not_nil",
        ]:
            self._global.declare(name)

    def check(self, program: Program) -> List[TypeDiag]:
        self.diags = []
        self._check_block(program.statements, self._global)
        return self.diags

    def _err(self, msg: str, line: int = 0):
        self.diags.append(TypeDiag("error", line, msg))

    def _warn(self, msg: str, line: int = 0):
        self.diags.append(TypeDiag("warning", line, msg))

    def _check_block(self, stmts: list, scope: TypeScope):
        for stmt in stmts:
            self._check_node(stmt, scope)

    def _check_node(self, node, scope: TypeScope):
        t = type(node).__name__

        if t == "FunctionDef":
            scope.declare_func(node.name, len(node.params))
            scope.declare(node.name)
            fn_scope = TypeScope(scope)
            # v0.2: register generic type params as type variables in the function scope
            for tp in getattr(node, "type_params", []):
                fn_scope.declare_type_var(tp)
            for p in node.params:
                ph = node.param_types.get(p)
                # If the annotation is a registered type-var, treat as Any for checking
                if ph and fn_scope.is_type_var(ph): ph = "Any"
                fn_scope.declare(p, ph)
            self._check_block(node.body, fn_scope)

        elif t == "ClassDef":
            scope.declare(node.name)
            cls_scope = TypeScope(scope)
            cls_scope.declare("self")
            # Interface conformance check
            for iface in node.implements:
                required = self._interfaces.get(iface, [])
                defined = [
                    n.name for n in node.body
                    if isinstance(n, FunctionDef)
                ]
                for req in required:
                    if req not in defined:
                        self._err(f"Class '{node.name}' missing method '{req}' required by interface '{iface}'")
            self._check_block(node.body, cls_scope)

        elif t == "InterfaceDef":
            self._interfaces[node.name] = node.method_names
            scope.declare(node.name)

        elif t == "LetStatement":
            hint = getattr(node, "type_hint", None)
            scope.declare(node.name, hint, is_let=True)   # v25: mark as let-bound
            expr_type = self._check_expr(node.value, scope)
            # v25: null safety
            if hint and not hint.endswith("?"):
                if expr_type in ("Null", "None") or self._is_null_literal(node.value):
                    self._err(
                        f"Null safety: cannot assign null to '{node.name}' of type '{hint}'. "
                        f"Use '{hint}?' to allow null.",
                        getattr(node, 'line', 0)
                    )
            if hint and expr_type and expr_type not in ("Any", "Null"):
                if not hint.endswith("?") and hint != expr_type and hint != "Any":
                    self._err(
                        f"Type mismatch: cannot assign {expr_type} to '{node.name}' of type {hint}",
                        getattr(node, 'line', 0)
                    )
        elif t == "VarStatement":    # v25: var x: Type = expr
            hint = getattr(node, "type_hint", None)
            scope.declare(node.name, hint if isinstance(hint, str) else None)
            expr_type = self._check_expr(node.value, scope)
            if hint and isinstance(hint, str) and expr_type and expr_type not in ("Any", "Null"):
                if not hint.endswith("?") and hint != expr_type and hint != "Any":
                    self._err(
                        f"Type mismatch: cannot assign {expr_type} to var '{node.name}' of type {hint}",
                        getattr(node, 'line', 0)
                    )

        elif t == "Assignment":
            declared_type = scope.lookup(node.name)
            if declared_type is False:
                self._err(f"Variable '{node.name}' used before declaration",
                          getattr(node, 'line', 0))
            # v25: warn on let-binding reassignment
            if scope.is_let_bound(node.name):
                self._warn(
                    f"Reassigning let-bound '{node.name}'. Consider using 'var' for mutable bindings.",
                    getattr(node, 'line', 0)
                )
            expr_type = self._check_expr(node.value, scope)
            if declared_type and expr_type and expr_type != "Any" and declared_type != "Any":
                if declared_type != expr_type:
                    self._err(f"Type mismatch: cannot assign {expr_type} to '{node.name}' of type {declared_type}",
                              getattr(node, 'line', 0))

        elif t in ("IfExpr", "WhileStatement"):
            self._check_expr(node.condition, scope)
            if t == "IfExpr":
                self._check_block(node.then_block, TypeScope(scope))
                self._check_block(node.else_block, TypeScope(scope))
            else:
                self._check_block(node.body, TypeScope(scope))

        elif t == "ForStatement":
            self._check_expr(node.iterable, scope)
            loop_scope = TypeScope(scope)
            loop_scope.declare(node.var)
            self._check_block(node.body, loop_scope)

        elif t == "ReturnStatement":
            if node.value: self._check_expr(node.value, scope)

        elif t == "ThrowStatement":
            self._check_expr(node.value, scope)

        elif t in ("ExprStatement", "TestStatement"):
            self._check_expr(getattr(node, "expr", getattr(node, "body", None)), scope)

        elif t == "TryExpr":
            self._check_block(node.try_block, TypeScope(scope))
            if node.catch_block:
                c_scope = TypeScope(scope)
                if node.catch_bind: c_scope.declare(node.catch_bind)
                self._check_block(node.catch_block, c_scope)
            if node.finally_block:
                self._check_block(node.finally_block, TypeScope(scope))

        elif t == "BlockExpr":
            self._check_block(node.stmts, TypeScope(scope))

        elif t == "NativeAssertStatement":   # v25: assert expr
            self._check_expr(node.condition, scope)
            if node.message is not None:
                self._check_expr(node.message, scope)

        elif t == "TraitDef":                # v25
            scope._traits[node.name] = {"methods": [m.name for m in node.methods]}
            scope.declare(node.name)
            for m in node.methods:
                if m.body:
                    fn_scope = TypeScope(scope)
                    for p in m.params:
                        fn_scope.declare(p)
                    self._check_block(m.body, fn_scope)

        elif t == "TraitImpl":               # v25: check all trait required methods are implemented
            trait = scope._traits.get(node.trait_name)
            if trait:
                required = trait.get("methods", [])
                provided = [m.name for m in node.methods]
                for req in required:
                    if req not in provided:
                        self._warn(
                            f"impl '{node.trait_name}' for '{node.class_name}' is missing method '{req}'"
                        )
            impl_scope = TypeScope(scope)
            impl_scope.declare("self")
            self._check_block(node.methods, impl_scope)

        elif t == "UseStatement":            # v25: use std::math::sin
            if node.names:
                for n in node.names:
                    scope.declare(n)
            elif node.path:
                scope.declare(node.alias or node.path[-1])

        # else: silently skip unknown nodes (expressions used as statements)

    def _is_null_literal(self, node) -> bool:
        """Check if a node is a NullLiteral in the AST."""
        return type(node).__name__ in ("NullLiteral", "NullExpr")

    def _check_expr(self, node, scope: TypeScope) -> str:
        if node is None: return "Any"
        t = type(node).__name__

        if t == "NumberLiteral":
            return "Float" if isinstance(node.value, float) else "Int"
        elif t == "StringLiteral": return "String"
        elif t == "BoolLiteral": return "Bool"
        elif t == "NullLiteral": return "Null"
        elif t in ("ListLiteral", "ArrayLiteral"):
            for el in getattr(node, "elements", []): self._check_expr(el, scope)
            return "List"
        elif t == "DictLiteral":
            for k, v in node.pairs: self._check_expr(k, scope); self._check_expr(v, scope)
            return "Dict"

        if t == "Identifier":
            declared = scope.lookup(node.name)
            if declared is False:
                self._err(f"Undefined variable: '{node.name}'")
                return "Any"
            return declared if isinstance(declared, str) else "Any"

        elif t == "FunctionCall":
            fn_name = node.callee.name if isinstance(node.callee, Identifier) else None
            if fn_name and not scope.lookup(fn_name):
                self._err(f"Undefined function: '{fn_name}'")
            expected = scope.func_arity(fn_name) if fn_name else None
            if expected is not None and len(node.args) != expected:
                if expected > 0 and len(node.args) == 0:
                    self._warn(f"Function '{fn_name}' expects {expected} arg(s), got 0")
            for arg in node.args:
                self._check_expr(arg, scope)

        elif t == "BinaryOp":
            self._check_expr(node.left, scope)
            self._check_expr(node.right, scope)

        elif t == "RangeExpr":
            self._check_expr(node.start, scope)
            self._check_expr(node.end, scope)

        elif t in ("MemberAccess", "OptionalAccess"):
            self._check_expr(node.obj, scope)

        elif t == "IndexAccess":
            self._check_expr(node.obj, scope)
            self._check_expr(node.index, scope)

        elif t in ("ListLiteral", "ArrayLiteral"):
            for el in node.elements: self._check_expr(el, scope)

        elif t == "DictLiteral":
            for k, v in node.pairs: self._check_expr(k, scope); self._check_expr(v, scope)

        elif t == "LambdaExpr":
            lam_scope = TypeScope(scope)
            for p in node.params: lam_scope.declare(p)
            self._check_expr(node.body, lam_scope)

        elif t == "UnaryOp":
            self._check_expr(node.operand, scope)

        elif t in ("TernaryOp", "TernaryExpr"):    # TernaryExpr is the v0.2 backward-compat alias
            self._check_expr(node.condition, scope)
            self._check_expr(node.true_expr, scope)
            self._check_expr(node.false_expr, scope)

        elif t == "GenericTypeAnnotation":           # v0.2: List<Int>, Dict<String, Int>
            # base type must be a declared name
            if not scope.lookup(node.base):
                self._err(f"Undefined generic type base: '{node.base}'")
            # each param should either be a type-var or a known type name
            for param in node.params:
                if not scope.is_type_var(param) and param not in (
                    "Int", "Float", "String", "Bool", "Any", "Null", "None"
                ):
                    self._warn(
                        f"Generic type param '{param}' is not a known type or type variable. "
                        f"Declare it explicitly with 'type {param} = ...' or a type-var."
                    )
            return f"{node.base}<{', '.join(node.params)}>"

        elif t == "GenericRuntimeCheck":             # v0.2: warn that this is runtime-only enforcement
            self._warn(
                f"GenericRuntimeCheck on '{getattr(node, 'item_type', '?')}': "
                f"This is a runtime assertion, NOT a compile-time type guarantee. "
                f"Use GenericTypeAnnotation in function signatures for static checking."
            )
            self._check_expr(node.target, scope)

        return "Any"
