const std = @import("std");
const parser_mod = @import("parser.zig");
const Program = parser_mod.Program;
const FnDecl = parser_mod.FnDecl;
const Stmt = parser_mod.Stmt;
const Expr = parser_mod.Expr;
const DataType = parser_mod.DataType;

pub const TypeError = error{
    TypeMismatch,
    UndefinedVariable,
    RedeclaredVariable,
    InvalidFieldAccess,
    InvalidFunctionCall,
    InvalidOptionalUnwrap,
    OutOfMemory,
};

pub const Symbol = struct {
    name: []const u8,
    data_type: DataType,
    is_const: bool,
};

pub const Scope = struct {
    parent: ?*Scope,
    symbols: std.StringHashMap(Symbol),

    pub fn init(allocator: std.mem.Allocator, parent: ?*Scope) Scope {
        return .{ .parent = parent, .symbols = std.StringHashMap(Symbol).init(allocator) };
    }

    pub fn deinit(self: *Scope) void {
        self.symbols.deinit();
    }

    pub fn put(self: *Scope, symbol: Symbol) TypeError!void {
        if (self.symbols.contains(symbol.name)) {
            return TypeError.RedeclaredVariable;
        }
        self.symbols.put(symbol.name, symbol) catch return TypeError.OutOfMemory;
    }

    pub fn get(self: *Scope, name: []const u8) ?Symbol {
        if (self.symbols.get(name)) |sym| return sym;
        if (self.parent) |p| return p.get(name);
        return null;
    }
};

/// Static scope/type validator. NOTE: since the pivot to a general-purpose
/// language (see PIVOT.md), `rinne run` uses the bytecode compiler directly;
/// the checker remains available for tooling and future phases.
pub const TypeChecker = struct {
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator) TypeChecker {
        return .{ .allocator = allocator };
    }

    pub fn checkProgram(self: *TypeChecker, program: *const Program) TypeError!void {
        var global_scope = Scope.init(self.allocator, null);
        defer global_scope.deinit();

        for (program.fn_decls) |fn_decl| {
            try global_scope.put(.{ .name = fn_decl.name, .data_type = .{ .primitive = .void_type }, .is_const = true });
        }

        for (program.global_stmts.stmts) |*stmt| {
            try self.checkStmt(&global_scope, stmt);
        }

        for (program.fn_decls) |*fn_decl| {
            try self.checkFnDecl(&global_scope, fn_decl);
        }
    }

    fn checkFnDecl(self: *TypeChecker, parent_scope: *Scope, fn_decl: *const FnDecl) TypeError!void {
        var fn_scope = Scope.init(self.allocator, parent_scope);
        defer fn_scope.deinit();

        for (fn_decl.params) |param| {
            try fn_scope.put(.{ .name = param.name, .data_type = param.param_type, .is_const = true });
        }

        for (fn_decl.body.stmts) |*stmt| {
            try self.checkStmt(&fn_scope, stmt);
        }
    }

    fn checkStmt(self: *TypeChecker, scope: *Scope, stmt: *const Stmt) TypeError!void {
        switch (stmt.*) {
            .var_decl => |decl| {
                const init_type = try self.checkExpr(scope, &decl.initializer);
                var final_type = init_type;
                if (decl.type_ann) |ann| {
                    if (!self.areTypesCompatible(ann, init_type)) {
                        return TypeError.TypeMismatch;
                    }
                    final_type = ann;
                }
                try scope.put(.{ .name = decl.name, .data_type = final_type, .is_const = decl.is_const });
            },
            .assignment => |assign| {
                const val_type = try self.checkExpr(scope, &assign.value);
                if (assign.target == .identifier) {
                    if (scope.get(assign.target.identifier)) |sym| {
                        if (sym.is_const) return TypeError.TypeMismatch;
                        if (!self.areTypesCompatible(sym.data_type, val_type)) {
                            return TypeError.TypeMismatch;
                        }
                    } else {
                        return TypeError.UndefinedVariable;
                    }
                } else {
                    _ = try self.checkExpr(scope, &assign.target);
                }
            },
            .expr_stmt => |*expr| {
                _ = try self.checkExpr(scope, expr);
            },
            .if_stmt => |if_s| {
                const cond_type = try self.checkExpr(scope, &if_s.condition);
                if (cond_type != .primitive or cond_type.primitive != .bool) {
                    return TypeError.TypeMismatch;
                }
                var then_scope = Scope.init(self.allocator, scope);
                defer then_scope.deinit();
                for (if_s.then_branch.stmts) |*s| {
                    try self.checkStmt(&then_scope, s);
                }
                if (if_s.else_branch) |eb| {
                    var else_scope = Scope.init(self.allocator, scope);
                    defer else_scope.deinit();
                    for (eb.stmts) |*s| {
                        try self.checkStmt(&else_scope, s);
                    }
                }
            },
            .while_stmt => |w| {
                const cond_type = try self.checkExpr(scope, &w.condition);
                if (cond_type != .primitive or cond_type.primitive != .bool) {
                    return TypeError.TypeMismatch;
                }
                var body_scope = Scope.init(self.allocator, scope);
                defer body_scope.deinit();
                for (w.body.stmts) |*s| {
                    try self.checkStmt(&body_scope, s);
                }
            },
            .break_stmt, .continue_stmt => {},
            .for_stmt => |f| {
                _ = try self.checkExpr(scope, &f.iter);
                var loop_scope = Scope.init(self.allocator, scope);
                defer loop_scope.deinit();
                for (f.body.stmts) |*s| {
                    try self.checkStmt(&loop_scope, s);
                }
            },
            .range_for => |r| {
                _ = try self.checkExpr(scope, &r.lo);
                _ = try self.checkExpr(scope, &r.hi);
                var loop_scope = Scope.init(self.allocator, scope);
                defer loop_scope.deinit();
                for (r.body.stmts) |*s| {
                    try self.checkStmt(&loop_scope, s);
                }
            },
            .return_stmt => |ret_expr| {
                if (ret_expr) |*e| {
                    _ = try self.checkExpr(scope, e);
                }
            },
            .defer_stmt => |*e| {
                _ = try self.checkExpr(scope, e);
            },
            .fn_decl => |*fd| {
                try scope.put(.{ .name = fd.name, .data_type = .{ .primitive = .void_type }, .is_const = true });
                try self.checkFnDecl(scope, fd);
            },
        }
    }

    pub fn checkExpr(self: *TypeChecker, scope: *Scope, expr: *const Expr) TypeError!DataType {
        switch (expr.*) {
            .int_lit => return .{ .primitive = .int },
            .float_lit => return .{ .primitive = .float },
            .string_lit => return .{ .primitive = .str },
            .bool_lit => return .{ .primitive = .bool },
            .null_lit => return .{ .primitive = .inferred },
            .identifier => |name| {
                if (scope.get(name)) |sym| {
                    return sym.data_type;
                }
                return TypeError.UndefinedVariable;
            },
            .array_lit => |items| {
                if (items.len == 0) {
                    const elem_box = try self.allocator.create(DataType);
                    elem_box.* = .{ .primitive = .inferred };
                    return .{ .slice = elem_box };
                }
                const first_type = try self.checkExpr(scope, &items[0]);
                for (items[1..]) |*item| {
                    const it = try self.checkExpr(scope, item);
                    if (!self.areTypesCompatible(first_type, it)) {
                        return TypeError.TypeMismatch;
                    }
                }
                const elem_box = try self.allocator.create(DataType);
                elem_box.* = first_type;
                return .{ .slice = elem_box };
            },
            .struct_lit => return .{ .primitive = .inferred },
            .binary => |bin| {
                const lhs_t = try self.checkExpr(scope, bin.lhs);
                const rhs_t = try self.checkExpr(scope, bin.rhs);
                if (!self.areTypesCompatible(lhs_t, rhs_t)) {
                    return TypeError.TypeMismatch;
                }
                switch (bin.op) {
                    .equal_equal, .not_equal, .less, .less_equal, .greater, .greater_equal, .kw_and, .kw_or => return .{ .primitive = .bool },
                    else => return lhs_t,
                }
            },
            .unary => |un| {
                if (un.op == .kw_not) return .{ .primitive = .bool };
                return try self.checkExpr(scope, un.operand);
            },
            .field_access => |fa| {
                _ = try self.checkExpr(scope, fa.object);
                return .{ .primitive = .inferred };
            },
            .index_access => |ia| {
                const obj_type = try self.checkExpr(scope, ia.object);
                _ = try self.checkExpr(scope, ia.index);
                if (obj_type == .slice) {
                    return obj_type.slice.*;
                }
                return .{ .primitive = .inferred };
            },
            .call => |c| {
                _ = try self.checkExpr(scope, c.callee);
                for (c.args) |*arg| {
                    _ = try self.checkExpr(scope, arg);
                }
                return .{ .primitive = .inferred };
            },
            .call_directive => |cd| {
                // Directive intrinsics are backend-resolved (@stdout, @env,
                // @tensor, ...). Arguments are still validated.
                for (cd.args) |*arg| {
                    _ = try self.checkExpr(scope, arg);
                }
                return .{ .primitive = .inferred };
            },
            .enum_lit => return .{ .primitive = .inferred },
            .fn_lit => return .{ .primitive = .inferred },
        }
    }

    fn areTypesCompatible(self: *TypeChecker, expected: DataType, actual: DataType) bool {
        if (expected == .primitive and expected.primitive == .inferred) return true;
        if (actual == .primitive and actual.primitive == .inferred) return true;

        switch (expected) {
            .primitive => |p1| {
                if (actual == .primitive) {
                    return p1 == actual.primitive;
                }
                return false;
            },
            .optional => |opt1| {
                if (actual == .optional) {
                    return self.areTypesCompatible(opt1.*, actual.optional.*);
                }
                return self.areTypesCompatible(opt1.*, actual);
            },
            .slice => |sl1| {
                if (actual == .slice) {
                    return self.areTypesCompatible(sl1.*, actual.slice.*);
                }
                return false;
            },
            .custom => return true,
        }
    }
};

test "TypeChecker validates variable typing" {
    var tc = TypeChecker.init(std.testing.allocator);
    var scope = Scope.init(std.testing.allocator, null);
    defer scope.deinit();

    const expr = Expr{ .int_lit = 42 };
    const t = try tc.checkExpr(&scope, &expr);
    try std.testing.expectEqual(DataType.Primitive.int, t.primitive);
}