"""Type system implementation for YARA semantic validation."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TYPE_CHECKING

from yaraast.visitor import ASTVisitor

if TYPE_CHECKING:
    from yaraast.ast.base import YaraFile
    from yaraast.ast.conditions import (
        AtExpression,
        ForExpression,
        ForOfExpression,
        InExpression,
        OfExpression,
    )
    from yaraast.ast.expressions import (
        ArrayAccess,
        BinaryExpression,
        BooleanLiteral,
        DoubleLiteral,
        Expression,
        FunctionCall,
        Identifier,
        IntegerLiteral,
        MemberAccess,
        ParenthesesExpression,
        RangeExpression,
        RegexLiteral,
        SetExpression,
        StringCount,
        StringIdentifier,
        StringLength,
        StringLiteral,
        StringOffset,
        UnaryExpression,
    )
    from yaraast.ast.rules import Import, Rule


class YaraType(ABC):
    """Base class for YARA types."""

    @abstractmethod
    def __str__(self) -> str:
        """String representation of the type."""

    @abstractmethod
    def is_compatible_with(self, other: YaraType) -> bool:
        """Check if this type is compatible with another."""

    def is_numeric(self) -> bool:
        """Check if this is a numeric type."""
        return False

    def is_string_like(self) -> bool:
        """Check if this is a string-like type."""
        return False


# Add static type instances after the concrete classes are defined
def _init_static_types() -> None:
    """Initialize static type instances on YaraType class."""
    YaraType.INTEGER = IntegerType()
    YaraType.STRING = StringType()
    YaraType.BOOLEAN = BooleanType()
    YaraType.DOUBLE = DoubleType()
    YaraType.REGEX = RegexType()
    YaraType.UNKNOWN = UnknownType()


@dataclass
class IntegerType(YaraType):
    """Integer type."""

    def __str__(self) -> str:
        return "integer"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, IntegerType | DoubleType)

    def is_numeric(self) -> bool:
        return True


@dataclass
class DoubleType(YaraType):
    """Double/float type."""

    def __str__(self) -> str:
        return "double"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, IntegerType | DoubleType)

    def is_numeric(self) -> bool:
        return True


@dataclass
class StringType(YaraType):
    """String type."""

    def __str__(self) -> str:
        return "string"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, StringType)

    def is_string_like(self) -> bool:
        return True


@dataclass
class BooleanType(YaraType):
    """Boolean type."""

    def __str__(self) -> str:
        return "boolean"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, BooleanType)


@dataclass
class StringSetType(YaraType):
    """String set type (for string identifiers)."""

    def __str__(self) -> str:
        return "string_set"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, StringSetType)


@dataclass
class RangeType(YaraType):
    """Range type."""

    def __str__(self) -> str:
        return "range"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, RangeType)


@dataclass
class ModuleType(YaraType):
    """Module type with attributes."""

    module_name: str
    attributes: dict[str, YaraType]

    def __str__(self) -> str:
        return f"module({self.module_name})"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, ModuleType) and other.module_name == self.module_name

    def get_attribute_type(self, attr: str) -> YaraType | None:
        """Get type of module attribute."""
        return self.attributes.get(attr)


@dataclass
class ArrayType(YaraType):
    """Array type."""

    element_type: YaraType

    def __str__(self) -> str:
        return f"array[{self.element_type}]"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, ArrayType) and self.element_type.is_compatible_with(
            other.element_type,
        )


@dataclass
class DictionaryType(YaraType):
    """Dictionary type."""

    key_type: YaraType
    value_type: YaraType

    def __str__(self) -> str:
        return f"dict[{self.key_type}, {self.value_type}]"

    def is_compatible_with(self, other: YaraType) -> bool:
        return (
            isinstance(other, DictionaryType)
            and self.key_type.is_compatible_with(other.key_type)
            and self.value_type.is_compatible_with(other.value_type)
        )


@dataclass
class FunctionType(YaraType):
    """Function type."""

    name: str
    param_types: list[YaraType]
    return_type: YaraType

    def __str__(self) -> str:
        params = ", ".join(str(p) for p in self.param_types)
        return f"{self.name}({params}) -> {self.return_type}"

    def is_compatible_with(self, other: YaraType) -> bool:
        return False  # Functions are not directly comparable


@dataclass
class UnknownType(YaraType):
    """Unknown type (for unresolved references)."""

    def __str__(self) -> str:
        return "unknown"

    def is_compatible_with(self, other: YaraType) -> bool:
        return True  # Unknown is compatible with anything


@dataclass
class RegexType(YaraType):
    """Regex type."""

    def __str__(self) -> str:
        return "regex"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, RegexType | StringType)

    def is_string_like(self) -> bool:
        return True


@dataclass
class StringIdentifierType(YaraType):
    """String identifier type ($a, $b, etc.)."""

    def __str__(self) -> str:
        return "string_identifier"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(
            other,
            StringType | RegexType | StringIdentifierType | BooleanType,
        )

    def is_string_like(self) -> bool:
        return True


@dataclass
class AnyType(YaraType):
    """Any type (variable or unspecified)."""

    def __str__(self) -> str:
        return "any"

    def is_compatible_with(self, other: YaraType) -> bool:
        return True


@dataclass
class FloatType(YaraType):
    """Float type (alias for double)."""

    def __str__(self) -> str:
        return "float"

    def is_compatible_with(self, other: YaraType) -> bool:
        return isinstance(other, FloatType | DoubleType | IntegerType)

    def is_numeric(self) -> bool:
        return True


@dataclass
class StructType(YaraType):
    """Struct type with named fields."""

    fields: dict[str, YaraType] = field(default_factory=dict)

    def __str__(self) -> str:
        return f"struct({', '.join(f'{k}: {v}' for k, v in self.fields.items())})"

    def is_compatible_with(self, other: YaraType) -> bool:
        if not isinstance(other, StructType):
            return False
        # Check all fields are compatible
        for field_name, field_type in self.fields.items():
            if field_name not in other.fields:
                return False
            if not field_type.is_compatible_with(other.fields[field_name]):
                return False
        return True


@dataclass
class FunctionDefinition:
    """Definition of a module function."""

    name: str
    return_type: YaraType
    parameters: list[tuple[str, YaraType]] = field(default_factory=list)


@dataclass
class ModuleDefinition:
    """Definition of a YARA module."""

    name: str
    attributes: dict[str, YaraType] = field(default_factory=dict)
    functions: dict[str, FunctionDefinition] = field(default_factory=dict)
    constants: dict[str, YaraType] = field(default_factory=dict)


class TypeSystem:
    """Type system with module support."""

    def __init__(self) -> None:
        self.modules: dict[str, ModuleDefinition] = {}
        self._init_modules()

    def _init_modules(self) -> None:
        """Initialize modules using ModuleLoader."""
        try:
            from yaraast.types.module_loader import ModuleLoader

            loader = ModuleLoader()
            self.modules = loader.modules
        except ImportError:
            # Fallback to hardcoded modules
            self._init_builtin_modules()

    def _init_builtin_modules(self) -> None:
        """Initialize builtin modules (fallback)."""
        # PE module
        pe = ModuleDefinition(name="pe")
        pe.attributes = {
            "machine": IntegerType(),
            "number_of_sections": IntegerType(),
            "timestamp": IntegerType(),
            "characteristics": IntegerType(),
            "entry_point": IntegerType(),
            "image_base": IntegerType(),
            "sections": ArrayType(
                StructType(
                    {
                        "name": StringType(),
                        "virtual_address": IntegerType(),
                        "virtual_size": IntegerType(),
                        "raw_size": IntegerType(),
                        "characteristics": IntegerType(),
                    },
                ),
            ),
            "version_info": DictionaryType(StringType(), StringType()),
            "number_of_resources": IntegerType(),
            "resource_timestamp": IntegerType(),
            "imports": ArrayType(StringType()),
            "exports": ArrayType(StringType()),
            "is_pe": BooleanType(),
            "is_dll": BooleanType(),
            "is_32bit": BooleanType(),
            "is_64bit": BooleanType(),
        }
        pe.functions = {
            "imphash": FunctionDefinition("imphash", StringType()),
            "section_index": FunctionDefinition(
                "section_index",
                IntegerType(),
                [("name", StringType())],
            ),
            "exports": FunctionDefinition(
                "exports",
                BooleanType(),
                [("name", StringType())],
            ),
            "imports": FunctionDefinition(
                "imports",
                BooleanType(),
                [("dll", StringType()), ("function", StringType())],
            ),
            "locale": FunctionDefinition(
                "locale",
                BooleanType(),
                [("locale", IntegerType())],
            ),
            "language": FunctionDefinition(
                "language",
                BooleanType(),
                [("lang", IntegerType())],
            ),
            # Add functions that can also be called as attributes
            "is_dll": FunctionDefinition("is_dll", BooleanType()),
            "is_64bit": FunctionDefinition("is_64bit", BooleanType()),
            "is_32bit": FunctionDefinition("is_32bit", BooleanType()),
            "rva_to_offset": FunctionDefinition(
                "rva_to_offset",
                IntegerType(),
                [("rva", IntegerType())],
            ),
        }
        self.modules["pe"] = pe

        # Math module
        math = ModuleDefinition(name="math")
        math.functions = {
            "abs": FunctionDefinition("abs", IntegerType(), [("x", IntegerType())]),
            "min": FunctionDefinition(
                "min",
                IntegerType(),
                [("a", IntegerType()), ("b", IntegerType())],
            ),
            "max": FunctionDefinition(
                "max",
                IntegerType(),
                [("a", IntegerType()), ("b", IntegerType())],
            ),
            "to_string": FunctionDefinition(
                "to_string",
                StringType(),
                [("n", IntegerType()), ("base", IntegerType())],
            ),
            "to_number": FunctionDefinition(
                "to_number",
                IntegerType(),
                [("s", StringType())],
            ),
            "log": FunctionDefinition("log", DoubleType(), [("x", DoubleType())]),
            "log2": FunctionDefinition("log2", DoubleType(), [("x", DoubleType())]),
            "log10": FunctionDefinition("log10", DoubleType(), [("x", DoubleType())]),
            "sqrt": FunctionDefinition("sqrt", DoubleType(), [("x", DoubleType())]),
        }
        self.modules["math"] = math

        # Add more modules as needed...

    def get_module(self, name: str) -> ModuleDefinition | None:
        """Get module definition by name."""
        return self.modules.get(name)


# Module definitions (deprecated - use TypeSystem instead)
MODULE_DEFINITIONS = {
    "pe": ModuleType(
        "pe",
        {
            "machine": IntegerType(),
            "number_of_sections": IntegerType(),
            "timestamp": IntegerType(),
            "characteristics": IntegerType(),
            "entry_point": IntegerType(),
            "image_base": IntegerType(),
            "sections": ArrayType(DictionaryType(StringType(), IntegerType())),
            "version_info": DictionaryType(StringType(), StringType()),
            "exports": ArrayType(StringType()),
            "imports": ArrayType(DictionaryType(StringType(), ArrayType(StringType()))),
            "is_dll": BooleanType(),
            "is_32bit": BooleanType(),
            "is_64bit": BooleanType(),
        },
    ),
    "elf": ModuleType(
        "elf",
        {
            "type": IntegerType(),
            "machine": IntegerType(),
            "entry_point": IntegerType(),
            "sections": ArrayType(DictionaryType(StringType(), IntegerType())),
            "segments": ArrayType(DictionaryType(StringType(), IntegerType())),
        },
    ),
    "math": ModuleType(
        "math",
        {
            "entropy": FunctionType(
                "entropy",
                [IntegerType(), IntegerType()],
                DoubleType(),
            ),
            "serial_correlation": FunctionType(
                "serial_correlation",
                [IntegerType(), IntegerType()],
                DoubleType(),
            ),
            "monte_carlo_pi": FunctionType(
                "monte_carlo_pi",
                [IntegerType(), IntegerType()],
                DoubleType(),
            ),
            "mean": FunctionType("mean", [IntegerType(), IntegerType()], DoubleType()),
            "deviation": FunctionType(
                "deviation",
                [IntegerType(), IntegerType(), DoubleType()],
                DoubleType(),
            ),
        },
    ),
    "dotnet": ModuleType(
        "dotnet",
        {
            "version": StringType(),
            "module_name": StringType(),
            "assembly": DictionaryType(StringType(), StringType()),
            "resources": ArrayType(DictionaryType(StringType(), IntegerType())),
            "streams": ArrayType(DictionaryType(StringType(), IntegerType())),
        },
    ),
    "hash": ModuleType(
        "hash",
        {
            "md5": FunctionType("md5", [IntegerType(), IntegerType()], StringType()),
            "sha1": FunctionType("sha1", [IntegerType(), IntegerType()], StringType()),
            "sha256": FunctionType(
                "sha256",
                [IntegerType(), IntegerType()],
                StringType(),
            ),
            "crc32": FunctionType(
                "crc32",
                [IntegerType(), IntegerType()],
                IntegerType(),
            ),
        },
    ),
}


class TypeEnvironment:
    """Type environment for tracking variable types."""

    def __init__(self) -> None:
        self.scopes: list[dict[str, YaraType]] = [{}]
        self.modules: set[str] = set()
        self.module_aliases: dict[str, str] = {}  # alias -> actual module name
        self.strings: set[str] = set()
        self.rules: set[str] = set()  # Track rule names

    def push_scope(self) -> None:
        """Push a new scope."""
        self.scopes.append({})

    def pop_scope(self) -> None:
        """Pop the current scope."""
        if len(self.scopes) > 1:
            self.scopes.pop()

    def define(self, name: str, type: YaraType) -> None:
        """Define a variable in the current scope."""
        self.scopes[-1][name] = type

    def lookup(self, name: str) -> YaraType | None:
        """Look up a variable type."""
        for scope in reversed(self.scopes):
            if name in scope:
                return scope[name]
        return None

    def add_module(self, alias: str, module: str | None = None) -> None:
        """Add an imported module with optional alias."""
        if module is None:
            # No alias, just module name
            self.modules.add(alias)
        else:
            # With alias
            self.modules.add(module)
            self.module_aliases[alias] = module

    def add_string(self, string_id: str) -> None:
        """Add a string identifier."""
        self.strings.add(string_id)

    def has_module(self, name: str) -> bool:
        """Check if module is imported (by name or alias)."""
        # Check if it's a direct module name
        if name in self.modules:
            return True
        # Check if it's an alias
        return name in self.module_aliases

    def get_module_name(self, name: str) -> str | None:
        """Get actual module name from alias or name."""
        if name in self.module_aliases:
            return self.module_aliases[name]
        if name in self.modules:
            return name
        return None

    def has_string(self, string_id: str) -> bool:
        """Check if string is defined."""
        return string_id in self.strings

    def has_string_pattern(self, pattern: str) -> bool:
        """Check if string pattern matches any defined strings.

        Supports wildcard patterns like $str* which matches $str1, $str2, etc.
        """
        if not pattern.endswith("*"):
            return self.has_string(pattern)

        # Handle wildcard patterns
        prefix = pattern[:-1]  # Remove the asterisk
        return any(string_id.startswith(prefix) for string_id in self.strings)

    def add_rule(self, rule_name: str) -> None:
        """Add a rule name."""
        self.rules.add(rule_name)

    def has_rule(self, rule_name: str) -> bool:
        """Check if rule is defined."""
        return rule_name in self.rules


class TypeInference(ASTVisitor[YaraType]):
    """Type inference visitor for expressions."""

    def __init__(self, env: TypeEnvironment) -> None:
        self.env = env
        self.errors: list[str] = []

    def infer(self, node: Expression) -> YaraType:
        """Infer type of expression."""
        return self.visit(node)

    # Literals
    def visit_integer_literal(self, node: IntegerLiteral) -> YaraType:
        return IntegerType()

    def visit_double_literal(self, node: DoubleLiteral) -> YaraType:
        return DoubleType()

    def visit_string_literal(self, node: StringLiteral) -> YaraType:
        return StringType()

    def visit_regex_literal(self, node: RegexLiteral) -> YaraType:
        return RegexType()  # Regex literals have their own type

    def visit_boolean_literal(self, node: BooleanLiteral) -> YaraType:
        return BooleanType()

    # Identifiers
    def visit_identifier(self, node: Identifier) -> YaraType:
        if node.name in {"filesize", "entrypoint"}:
            return IntegerType()

        # VT LiveHunt legacy identifiers (work without import "vt")
        if node.name == "new_file":
            return BooleanType()
        if node.name in {"positives", "submissions"}:
            return IntegerType()

        if node.name == "them":
            return StringSetType()

        # Handle YARA quantifier keywords
        if node.name in ("any", "all", "none"):
            return StringType()  # Quantifiers are treated as string keywords

        # Check if it's a rule reference
        if self.env.has_rule(node.name):
            return BooleanType()  # Rule references evaluate to boolean

        # Check if it's a module name or alias
        if self.env.has_module(node.name):
            # Get the actual module name (handles aliases)
            actual_module = self.env.get_module_name(node.name)
            if actual_module:
                # Return module type based on module loader
                from yaraast.types.module_loader import ModuleLoader

                loader = ModuleLoader()
                module_def = loader.get_module(actual_module)
                if module_def:
                    # Convert ModuleDefinition to ModuleType
                    return ModuleType(
                        module_name=actual_module,
                        attributes=module_def.attributes,
                    )

        var_type = self.env.lookup(node.name)
        if var_type:
            return var_type

        return UnknownType()

    def visit_string_identifier(self, node: StringIdentifier) -> YaraType:
        if self.env.has_string(node.name) or self.env.has_string_pattern(node.name):
            # String identifiers are dual-typed: string-like for operators like 'matches'
            # but boolean-compatible for logical operations
            return StringIdentifierType()
        self.errors.append(f"Undefined string: {node.name}")
        return UnknownType()

    def visit_string_wildcard(self, node) -> YaraType:
        """Visit string wildcard ($*)."""
        return StringSetType()

    def visit_string_count(self, node: StringCount) -> YaraType:
        string_id = f"${node.string_id}" if not node.string_id.startswith("$") else node.string_id
        if self.env.has_string(string_id) or self.env.has_string_pattern(string_id):
            return IntegerType()
        self.errors.append(f"Undefined string: {string_id}")
        return UnknownType()

    def visit_string_offset(self, node: StringOffset) -> YaraType:
        string_id = f"${node.string_id}" if not node.string_id.startswith("$") else node.string_id
        if self.env.has_string(string_id) or self.env.has_string_pattern(string_id):
            # Additionally check index if present
            if hasattr(node, "index") and node.index:
                index_type = self.visit(node.index)
                if not isinstance(index_type, IntegerType):
                    self.errors.append(
                        f"String offset index must be integer, got {index_type}",
                    )
            return IntegerType()
        self.errors.append(f"Undefined string: {string_id}")
        return UnknownType()

    def visit_string_length(self, node: StringLength) -> YaraType:
        string_id = f"${node.string_id}" if not node.string_id.startswith("$") else node.string_id
        if self.env.has_string(string_id) or self.env.has_string_pattern(string_id):
            # Additionally check index if present
            if hasattr(node, "index") and node.index:
                index_type = self.visit(node.index)
                if not isinstance(index_type, IntegerType):
                    self.errors.append(
                        f"String length index must be integer, got {index_type}",
                    )
            return IntegerType()
        self.errors.append(f"Undefined string: {string_id}")
        return UnknownType()

    # Binary expressions
    def visit_binary_expression(self, node: BinaryExpression) -> YaraType:
        left_type = self.visit(node.left)
        right_type = self.visit(node.right)

        # Logical operators
        if node.operator in ["and", "or"]:
            if not isinstance(left_type, BooleanType | StringIdentifierType):
                self.errors.append(
                    f"Left operand of '{node.operator}' must be boolean, got {left_type}",
                )
            if not isinstance(right_type, BooleanType | StringIdentifierType):
                self.errors.append(
                    f"Right operand of '{node.operator}' must be boolean, got {right_type}",
                )
            return BooleanType()

        # Comparison operators
        if node.operator in ["<", "<=", ">", ">=", "==", "!="]:
            if (left_type.is_numeric() and right_type.is_numeric()) or isinstance(
                left_type,
                type(right_type),
            ):
                return BooleanType()
            self.errors.append(
                f"Incompatible types for '{node.operator}': {left_type} and {right_type}",
            )
            return BooleanType()

        # String operators
        if node.operator in [
            "contains",
            "matches",
            "startswith",
            "endswith",
            "icontains",
            "istartswith",
            "iendswith",
            "iequals",
        ]:
            # contains can also work on arrays (e.g., array contains element)
            if node.operator == "contains" and isinstance(left_type, ArrayType):
                # For arrays, right side should match element type
                if not left_type.element_type.is_compatible_with(right_type):
                    self.errors.append(
                        f"Array element type {left_type.element_type} not compatible with {right_type}",
                    )
                return BooleanType()

            if not left_type.is_string_like():
                self.errors.append(
                    f"Left operand of '{node.operator}' must be string-like or array, got {left_type}",
                )
            # For matches, right side can be string or regex
            if node.operator == "matches":
                if not isinstance(right_type, StringType | RegexType):
                    self.errors.append(
                        f"Right operand of 'matches' must be string or regex, got {right_type}",
                    )
            elif not isinstance(right_type, StringType):
                self.errors.append(
                    f"Right operand of '{node.operator}' must be string, got {right_type}",
                )
            return BooleanType()

        # Arithmetic operators
        if node.operator in ["+", "-", "*", "/", "%"]:
            if not left_type.is_numeric():
                self.errors.append(
                    f"Left operand of '{node.operator}' must be numeric, got {left_type}",
                )
            if not right_type.is_numeric():
                self.errors.append(
                    f"Right operand of '{node.operator}' must be numeric, got {right_type}",
                )

            # Division always returns double
            if (
                node.operator == "/"
                or isinstance(left_type, DoubleType)
                or isinstance(right_type, DoubleType)
            ):
                return DoubleType()
            return IntegerType()

        # Bitwise operators
        if node.operator in ["&", "|", "^", "<<", ">>"]:
            if not isinstance(left_type, IntegerType):
                self.errors.append(
                    f"Left operand of '{node.operator}' must be integer, got {left_type}",
                )
            if not isinstance(right_type, IntegerType):
                self.errors.append(
                    f"Right operand of '{node.operator}' must be integer, got {right_type}",
                )
            return IntegerType()

        self.errors.append(f"Unknown binary operator: {node.operator}")
        return UnknownType()

    def visit_unary_expression(self, node: UnaryExpression) -> YaraType:
        operand_type = self.visit(node.operand)

        if node.operator == "not":
            if not isinstance(operand_type, BooleanType):
                self.errors.append(
                    f"Operand of 'not' must be boolean, got {operand_type}",
                )
            return BooleanType()
        if node.operator == "-":
            if not operand_type.is_numeric():
                self.errors.append(
                    f"Operand of '-' must be numeric, got {operand_type}",
                )
            return operand_type
        if node.operator == "~":
            if not isinstance(operand_type, IntegerType):
                self.errors.append(
                    f"Operand of '~' must be integer, got {operand_type}",
                )
            return IntegerType()
        self.errors.append(f"Unknown unary operator: {node.operator}")
        return UnknownType()

    def visit_parentheses_expression(self, node: ParenthesesExpression) -> YaraType:
        return self.visit(node.expression)

    def visit_set_expression(self, node: SetExpression) -> YaraType:
        # Check all elements have same type
        if node.elements:
            first_type = self.visit(node.elements[0])
            for elem in node.elements[1:]:
                elem_type = self.visit(elem)
                if not first_type.is_compatible_with(elem_type):
                    self.errors.append(
                        f"Set elements must have same type: {first_type} vs {elem_type}",
                    )

        return StringSetType()  # Sets are typically string sets in YARA

    def visit_range_expression(self, node: RangeExpression) -> YaraType:
        low_type = self.visit(node.low)
        high_type = self.visit(node.high)

        if not isinstance(low_type, IntegerType):
            self.errors.append(f"Range low bound must be integer, got {low_type}")
        if not isinstance(high_type, IntegerType):
            self.errors.append(f"Range high bound must be integer, got {high_type}")

        return RangeType()

    def visit_function_call(self, node: FunctionCall) -> YaraType:
        # Check if it's a module function call (e.g., m.entropy)
        if "." in node.function:
            parts = node.function.split(".", 1)
            if len(parts) == 2:
                module_name, func_name = parts
                if self.env.has_module(module_name):
                    # Get the actual module name (handles aliases)
                    actual_module = self.env.get_module_name(module_name)
                    if actual_module:
                        from yaraast.types.module_loader import ModuleLoader

                        loader = ModuleLoader()
                        module_def = loader.get_module(actual_module)
                        if module_def and func_name in module_def.functions:
                            func_def = module_def.functions[func_name]
                            # Validate argument types
                            if func_def.parameters and len(node.arguments) != len(
                                func_def.parameters,
                            ):
                                self.errors.append(
                                    f"Function '{func_name}' expects {len(func_def.parameters)} arguments, got {len(node.arguments)}",
                                )
                            return func_def.return_type
                        self.errors.append(
                            f"Module '{actual_module}' has no function '{func_name}'",
                        )
                        return UnknownType()

        # Built-in functions
        if node.function == "uint8":
            if len(node.arguments) != 1:
                self.errors.append("uint8() expects 1 argument")
            return IntegerType()
        if node.function == "uint16":
            if len(node.arguments) != 1:
                self.errors.append("uint16() expects 1 argument")
            return IntegerType()
        if node.function == "uint32":
            if len(node.arguments) != 1:
                self.errors.append("uint32() expects 1 argument")
            return IntegerType()
        if node.function == "int8":
            if len(node.arguments) != 1:
                self.errors.append("int8() expects 1 argument")
            return IntegerType()
        if node.function == "int16":
            if len(node.arguments) != 1:
                self.errors.append("int16() expects 1 argument")
            return IntegerType()
        if node.function == "int32":
            if len(node.arguments) != 1:
                self.errors.append("int32() expects 1 argument")
            return IntegerType()
        # Big-endian variants
        if node.function == "uint8be":
            if len(node.arguments) != 1:
                self.errors.append("uint8be() expects 1 argument")
            return IntegerType()
        if node.function == "uint16be":
            if len(node.arguments) != 1:
                self.errors.append("uint16be() expects 1 argument")
            return IntegerType()
        if node.function == "uint32be":
            if len(node.arguments) != 1:
                self.errors.append("uint32be() expects 1 argument")
            return IntegerType()
        if node.function == "int8be":
            if len(node.arguments) != 1:
                self.errors.append("int8be() expects 1 argument")
            return IntegerType()
        if node.function == "int16be":
            if len(node.arguments) != 1:
                self.errors.append("int16be() expects 1 argument")
            return IntegerType()
        if node.function == "int32be":
            if len(node.arguments) != 1:
                self.errors.append("int32be() expects 1 argument")
            return IntegerType()
        # Little-endian variants (explicit)
        if node.function == "uint16le":
            if len(node.arguments) != 1:
                self.errors.append("uint16le() expects 1 argument")
            return IntegerType()
        if node.function == "uint32le":
            if len(node.arguments) != 1:
                self.errors.append("uint32le() expects 1 argument")
            return IntegerType()
        if node.function == "int16le":
            if len(node.arguments) != 1:
                self.errors.append("int16le() expects 1 argument")
            return IntegerType()
        if node.function == "int32le":
            if len(node.arguments) != 1:
                self.errors.append("int32le() expects 1 argument")
            return IntegerType()

        return UnknownType()

    def visit_array_access(self, node: ArrayAccess) -> YaraType:
        array_type = self.visit(node.array)
        index_type = self.visit(node.index)

        if not isinstance(index_type, IntegerType):
            self.errors.append(f"Array index must be integer, got {index_type}")

        if isinstance(array_type, ArrayType):
            return array_type.element_type
        self.errors.append(f"Cannot index non-array type: {array_type}")
        return UnknownType()

    def visit_member_access(self, node: MemberAccess) -> YaraType:
        obj_type = self.visit(node.object)

        if isinstance(obj_type, ModuleType):
            attr_type = obj_type.get_attribute_type(node.member)
            if attr_type:
                return attr_type
            self.errors.append(
                f"Module '{obj_type.module_name}' has no attribute '{node.member}'",
            )
            return UnknownType()

        if isinstance(obj_type, StructType):
            if node.member in obj_type.fields:
                return obj_type.fields[node.member]
            self.errors.append(
                f"Struct has no field '{node.member}'",
            )
            return UnknownType()

        self.errors.append(f"Cannot access member of non-module type: {obj_type}")
        return UnknownType()

    def visit_module_reference(self, node) -> YaraType:
        if self.env.has_module(node.module):
            # Get the actual module name (handles aliases)
            actual_module = self.env.get_module_name(node.module)
            if actual_module:
                # Return module type based on module loader
                from yaraast.types.module_loader import ModuleLoader

                loader = ModuleLoader()
                module_def = loader.get_module(actual_module)
                if module_def:
                    # Convert ModuleDefinition to ModuleType
                    return ModuleType(
                        module_name=actual_module,
                        attributes=module_def.attributes,
                    )
            # Fallback to old MODULE_DEFINITIONS for compatibility
            return MODULE_DEFINITIONS.get(node.module, UnknownType())
        self.errors.append(f"Module '{node.module}' not imported")
        return UnknownType()

    def visit_dictionary_access(self, node) -> YaraType:
        dict_type = self.visit(node.object)

        if isinstance(dict_type, DictionaryType):
            # Check key type if it's an expression
            if hasattr(node, "key") and hasattr(node.key, "accept"):
                key_type = self.visit(node.key)
                if not isinstance(key_type, dict_type.key_type.__class__):
                    self.errors.append(
                        f"Dictionary key must be {dict_type.key_type}, got {key_type}",
                    )
            return dict_type.value_type
        self.errors.append(f"Cannot access dictionary on non-dict type: {dict_type}")
        return UnknownType()

    # Conditions
    def visit_at_expression(self, node: AtExpression) -> YaraType:
        offset_type = self.visit(node.offset)
        if not isinstance(offset_type, IntegerType):
            self.errors.append(
                f"Offset in 'at' expression must be integer, got {offset_type}",
            )
        return BooleanType()

    def visit_in_expression(self, node: InExpression) -> YaraType:
        range_type = self.visit(node.range)
        if not isinstance(range_type, RangeType):
            self.errors.append(f"'in' expression requires range, got {range_type}")
        return BooleanType()

    def visit_of_expression(self, node: OfExpression) -> YaraType:
        # Quantifier should be string ("any", "all") or integer
        quant_type = self.visit(node.quantifier)
        if not isinstance(quant_type, StringType | IntegerType):
            self.errors.append(
                f"'of' quantifier must be string or integer, got {quant_type}",
            )

        # String set should be StringSetType
        set_type = self.visit(node.string_set)
        if not isinstance(set_type, StringSetType):
            self.errors.append(f"'of' requires string set, got {set_type}")

        return BooleanType()

    def visit_for_expression(self, node: ForExpression) -> YaraType:
        # Add loop variable to environment
        self.env.push_scope()

        # Infer iterable type
        iter_type = self.visit(node.iterable)

        # Determine loop variable type based on iterable
        if isinstance(iter_type, RangeType):
            self.env.define(node.variable, IntegerType())
        elif isinstance(iter_type, ArrayType):
            self.env.define(node.variable, iter_type.element_type)
        else:
            self.errors.append(f"Cannot iterate over type: {iter_type}")
            self.env.define(node.variable, UnknownType())

        # Check body returns boolean
        body_type = self.visit(node.body)
        if not isinstance(body_type, BooleanType):
            self.errors.append(f"For loop body must return boolean, got {body_type}")

        self.env.pop_scope()
        return BooleanType()

    def visit_for_of_expression(self, node: ForOfExpression) -> YaraType:
        # String set should be StringSetType
        set_type = self.visit(node.string_set)
        if not isinstance(set_type, StringSetType):
            self.errors.append(f"'for...of' requires string set, got {set_type}")

        # Condition should be boolean if present
        if node.condition:
            cond_type = self.visit(node.condition)
            if not isinstance(cond_type, BooleanType):
                self.errors.append(
                    f"'for...of' condition must be boolean, got {cond_type}",
                )

        return BooleanType()

    # Default implementations for other visit methods
    def visit_yara_file(self, node):
        return UnknownType()

    def visit_import(self, node):
        return UnknownType()

    def visit_include(self, node):
        return UnknownType()

    def visit_rule(self, node):
        return UnknownType()

    def visit_tag(self, node):
        return UnknownType()

    def visit_string_definition(self, node):
        return UnknownType()

    def visit_plain_string(self, node):
        return UnknownType()

    def visit_hex_string(self, node):
        return UnknownType()

    def visit_regex_string(self, node):
        return UnknownType()

    def visit_string_modifier(self, node):
        return UnknownType()

    def visit_hex_token(self, node):
        return UnknownType()

    def visit_hex_byte(self, node):
        return UnknownType()

    def visit_hex_wildcard(self, node):
        return UnknownType()

    def visit_hex_jump(self, node):
        return UnknownType()

    def visit_hex_alternative(self, node):
        return UnknownType()

    def visit_hex_nibble(self, node):
        return UnknownType()

    def visit_expression(self, node):
        return UnknownType()

    def visit_condition(self, node):
        return UnknownType()

    def visit_meta(self, node):
        return UnknownType()

    def visit_comment(self, node):
        return UnknownType()

    def visit_comment_group(self, node):
        return UnknownType()

    def visit_defined_expression(self, node):
        return BooleanType()

    def visit_string_operator_expression(self, node):
        return BooleanType()

    # Add missing abstract methods for TypeInference
    def visit_extern_import(self, node):
        return UnknownType()

    def visit_extern_namespace(self, node):
        return UnknownType()

    def visit_extern_rule(self, node):
        return UnknownType()

    def visit_extern_rule_reference(self, node):
        return UnknownType()

    def visit_in_rule_pragma(self, node):
        return UnknownType()

    def visit_pragma(self, node):
        return UnknownType()

    def visit_pragma_block(self, node):
        return UnknownType()


class TypeChecker(ASTVisitor[None]):
    """Type checker for YARA rules."""

    def __init__(self) -> None:
        self.env = TypeEnvironment()
        self.inference = TypeInference(self.env)
        self.errors: list[str] = []

    def check_compatibility(self, type1, type2) -> bool:
        """Check if two types are compatible.

        For compatibility with the test, we accept YaraType class attributes.
        """
        # Handle static YaraType attributes from the test
        if hasattr(type1, "__name__"):  # It's a class
            if type1.__name__ == "YaraType" and hasattr(type1, "INTEGER"):
                type1 = type1.INTEGER
            elif type1.__name__ == "YaraType" and hasattr(type1, "STRING"):
                type1 = type1.STRING

        if hasattr(type2, "__name__"):  # It's a class
            if type2.__name__ == "YaraType" and hasattr(type2, "INTEGER"):
                type2 = type2.INTEGER
            elif type2.__name__ == "YaraType" and hasattr(type2, "STRING"):
                type2 = type2.STRING

        # If they're instances, use the is_compatible_with method
        if isinstance(type1, YaraType) and isinstance(type2, YaraType):
            return type1.is_compatible_with(type2)

        # Otherwise check if they're the same
        return type1 == type2

    def infer_type(self, node):
        """Infer type from AST node."""
        return self.inference.infer(node)

    def check(self, ast: YaraFile) -> list[str]:
        """Type check a YARA file and return errors."""
        self.errors = []
        self.visit(ast)
        self.errors.extend(self.inference.errors)
        return self.errors

    def visit_yara_file(self, node: YaraFile) -> None:
        # Process imports first
        for imp in node.imports:
            self.visit(imp)

        # Add all rule names first to support forward references
        for rule in node.rules:
            self.env.add_rule(rule.name)

        # Process rules
        for rule in node.rules:
            self.visit(rule)

    def visit_import(self, node: Import) -> None:
        # Use alias if provided, otherwise use module name
        name = node.alias if node.alias else node.module
        self.env.add_module(name, node.module)

    def visit_rule(self, node: Rule) -> None:
        # Add string definitions to environment
        for string in node.strings:
            self.env.add_string(string.identifier)

        # Type check condition
        if node.condition:
            cond_type = self.inference.infer(node.condition)
            # In YARA, integer conditions are valid (0 = false, non-zero = true)
            # Also string counts and offsets return integers that can be used as conditions
            # String identifiers ($a, $b) are also valid as boolean conditions
            if not isinstance(
                cond_type,
                BooleanType | IntegerType | StringIdentifierType,
            ):
                self.errors.append(
                    f"Rule condition must be boolean, integer, or string identifier, got {cond_type}",
                )

    # Other visit methods with pass
    def visit_include(self, node) -> None:
        """Include directives don't affect type checking."""
        # Implementation intentionally empty

    def visit_tag(self, node) -> None:
        """Tags don't affect type checking."""
        # Implementation intentionally empty

    def visit_string_definition(self, node) -> None:
        """String definitions are handled at rule level."""
        # Implementation intentionally empty

    def visit_plain_string(self, node) -> None:
        """Plain strings are handled at rule level."""
        # Implementation intentionally empty

    def visit_hex_string(self, node) -> None:
        """Hex strings are handled at rule level."""
        # Implementation intentionally empty

    def visit_regex_string(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_modifier(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_token(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_byte(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_wildcard(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_jump(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_alternative(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_hex_nibble(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_identifier(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_identifier(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_wildcard(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_count(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_offset(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_length(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_integer_literal(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_double_literal(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_literal(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_boolean_literal(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_binary_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_unary_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_parentheses_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_set_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_range_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_function_call(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_array_access(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_member_access(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_condition(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_for_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_for_of_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_at_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_in_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_of_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_meta(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_module_reference(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_dictionary_access(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_comment(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_comment_group(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_defined_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_regex_literal(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_string_operator_expression(self, node) -> None:
        pass  # Implementation intentionally empty

    # Add missing abstract methods
    def visit_extern_import(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_extern_namespace(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_extern_rule(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_extern_rule_reference(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_in_rule_pragma(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_pragma(self, node) -> None:
        pass  # Implementation intentionally empty

    def visit_pragma_block(self, node) -> None:
        pass  # Implementation intentionally empty


class TypeValidator:
    """High-level type validation API."""

    @staticmethod
    def validate(ast: YaraFile) -> tuple[bool, list[str]]:
        """Validate types in YARA file. Returns (is_valid, errors)."""
        checker = TypeChecker()
        errors = checker.check(ast)
        return len(errors) == 0, errors

    @staticmethod
    def validate_expression(
        expr: Expression,
        env: TypeEnvironment | None = None,
    ) -> tuple[YaraType, list[str]]:
        """Validate and infer type of expression."""
        if env is None:
            env = TypeEnvironment()

        inference = TypeInference(env)
        expr_type = inference.infer(expr)
        return expr_type, inference.errors


# Initialize static type instances
_init_static_types()
