"""
Type Inference Engine - Milestone M1 (Extended in v0.3.0)
==========================================================

Infers Z3-compatible types from Python AST.

Scope:
- Integer literals and operations → INT
- Boolean literals and comparisons → BOOL
- String literals and operations → STRING (v0.3.0)
- Everything else → UNKNOWN

This is a flow-insensitive analysis: we track the LAST assignment to each variable.

# ==============================================
#
# COMMUNITY (Current & Planned):
#
# COMMUNITY Examples & Tutorials:
#
# COMMUNITY Testing & Validation:
#
# PRO (Enhanced Features):
#
# PRO Type System Features:
#
# ENTERPRISE (Advanced Capabilities):
"""

import ast
from enum import Enum, auto
from typing import Any

from z3 import BoolSort, IntSort, RealSort, StringSort

from code_scalpel.parsing.unified_parser import ParsingError, parse_python_code


class InferredType(Enum):
    """
    Types that can be inferred from Python code.

    Phase 1: INT and BOOL
    Phase 2 (v0.3.0): STRING added
    Phase 3 (v1.3.0): FLOAT added
    UNKNOWN means we cannot determine the type (unsupported construct).
    """

    INT = auto()
    BOOL = auto()
    STRING = auto()  # v0.3.0
    FLOAT = auto()  # v1.3.0
    UNKNOWN = auto()

    def to_z3_sort(self) -> Any:
        """Convert to Z3 sort. Raises if UNKNOWN."""
        if self == InferredType.INT:
            return IntSort()
        elif self == InferredType.BOOL:
            return BoolSort()
        elif self == InferredType.STRING:
            return StringSort()
        elif self == InferredType.FLOAT:
            return RealSort()  # Z3 uses Real for floats
        else:
            raise ValueError(f"Cannot convert {self.name} to Z3 sort")

    def __repr__(self) -> str:
        return f"InferredType.{self.name}"


class TypeInferenceEngine:
    """
    Infers types for variables in Python code.

    Usage:
        engine = TypeInferenceEngine()
        types = engine.infer("x = 1\\ny = x > 0")
        # types = {"x": InferredType.INT, "y": InferredType.BOOL}

    Design:
        - Flow-insensitive: tracks last assignment per variable
        - Conservative: unknown operations produce UNKNOWN
        - Taint propagation: UNKNOWN + anything = UNKNOWN
    """

    def __init__(self):
        self._types: dict[str, InferredType] = {}

    def infer(self, code: str) -> dict[str, InferredType]:
        """
        Infer types for all variables in the code.

        Args:
            code: Python source code as string

        Returns:
            Dict mapping variable names to their inferred types
        """
        self._types = {}

        if not code.strip():
            return {}

        try:
            tree, _report = parse_python_code(code, filename="<type_inference>")
        except ParsingError:
            return {}

        self._visit(tree)
        return self._types.copy()

    def _visit(self, node: ast.AST) -> None:
        """Visit AST node and its children."""
        if isinstance(node, ast.Assign):
            self._handle_assign(node)
        elif isinstance(node, ast.AugAssign):
            self._handle_aug_assign(node)

        # Visit children
        for child in ast.iter_child_nodes(node):
            self._visit(child)

    def _handle_assign(self, node: ast.Assign) -> None:
        """Handle assignment: x = expr or x = y = expr"""
        value_type = self._infer_expr_type(node.value)

        for target in node.targets:
            if isinstance(target, ast.Name):
                self._types[target.id] = value_type
            elif isinstance(target, ast.Tuple) or isinstance(target, ast.List):
                # Tuple unpacking: a, b = ... → both UNKNOWN
                for elt in target.elts:
                    if isinstance(elt, ast.Name):
                        self._types[elt.id] = InferredType.UNKNOWN

    def _handle_aug_assign(self, node: ast.AugAssign) -> None:
        """Handle augmented assignment: x += 1"""
        if isinstance(node.target, ast.Name):
            var_name = node.target.id
            current_type = self._types.get(var_name, InferredType.UNKNOWN)
            value_type = self._infer_expr_type(node.value)

            # Result type depends on operation and operands
            result_type = self._combine_types_for_binop(current_type, node.op, value_type)
            self._types[var_name] = result_type

    def _infer_expr_type(self, node: ast.expr) -> InferredType:
        """Infer the type of an expression."""

        # Integer literals
        if isinstance(
            node, ast.Constant
        ):  # [20251216_BUGFIX] Rely on Constant to avoid deprecated ast.Num/Str warnings
            return self._infer_constant_type(node.value)

        # Name reference: look up in our type map
        if isinstance(node, ast.Name):
            return self._types.get(node.id, InferredType.UNKNOWN)

        # Unary operations: -x, not x, +x, ~x
        if isinstance(node, ast.UnaryOp):
            return self._infer_unary_type(node)

        # Binary operations: x + y, x * y, etc.
        if isinstance(node, ast.BinOp):
            return self._infer_binop_type(node)

        # Boolean operations: x and y, x or y
        if isinstance(node, ast.BoolOp):
            return self._infer_boolop_type(node)

        # Comparisons: x < y, x == y, 1 < x < 10
        if isinstance(node, ast.Compare):
            return InferredType.BOOL

        # IfExp: x if cond else y
        if isinstance(node, ast.IfExp):
            # Both branches should have same type ideally
            # For now, return UNKNOWN (conservative)
            return InferredType.UNKNOWN

        # Everything else: function calls, attribute access, subscript, etc.
        # → UNKNOWN (Phase 1 limitation)
        return InferredType.UNKNOWN

    def _infer_constant_type(self, value) -> InferredType:
        """Infer type from a constant value."""
        if isinstance(value, bool):
            # IMPORTANT: bool check MUST come before int
            # because isinstance(True, int) is True in Python!
            return InferredType.BOOL
        elif isinstance(value, float):
            # v1.3.0: Float support - check before int since float is more specific
            return InferredType.FLOAT
        elif isinstance(value, int):
            return InferredType.INT
        elif isinstance(value, str):
            # v0.3.0: String support
            return InferredType.STRING
        else:
            # bytes, None, etc. → UNKNOWN
            return InferredType.UNKNOWN

    def _infer_unary_type(self, node: ast.UnaryOp) -> InferredType:
        """Infer type of unary operation."""
        operand_type = self._infer_expr_type(node.operand)

        if isinstance(node.op, ast.Not):
            # not x → Bool (regardless of operand in Python, but we're strict)
            return InferredType.BOOL

        elif isinstance(node.op, (ast.UAdd, ast.USub)):
            # +x, -x → same type as operand (if Int or Float, stays same)
            if operand_type == InferredType.INT:
                return InferredType.INT
            elif operand_type == InferredType.FLOAT:
                return InferredType.FLOAT
            return InferredType.UNKNOWN

        elif isinstance(node.op, ast.Invert):
            # ~x → Int (bitwise invert)
            if operand_type == InferredType.INT:
                return InferredType.INT
            return InferredType.UNKNOWN

        # All Python unary operators (Not, UAdd, USub, Invert) are handled above.
        # This fallback is defensive code for future Python versions.
        # If a new unary operator is added and this code is reached,
        # returning UNKNOWN is the safe default.
        return InferredType.UNKNOWN  # pragma: no cover

    def _infer_binop_type(self, node: ast.BinOp) -> InferredType:
        """Infer type of binary operation."""
        left_type = self._infer_expr_type(node.left)
        right_type = self._infer_expr_type(node.right)

        return self._combine_types_for_binop(left_type, node.op, right_type)

    def _combine_types_for_binop(self, left: InferredType, op: ast.operator, right: InferredType) -> InferredType:
        """Determine result type of binary operation."""

        # UNKNOWN taints everything
        if left == InferredType.UNKNOWN or right == InferredType.UNKNOWN:
            return InferredType.UNKNOWN

        # String concatenation: str + str → str
        if isinstance(op, ast.Add):
            if left == InferredType.STRING and right == InferredType.STRING:
                return InferredType.STRING
            if left == InferredType.INT and right == InferredType.INT:
                return InferredType.INT
            # v1.3.0: Float arithmetic
            if left == InferredType.FLOAT and right == InferredType.FLOAT:
                return InferredType.FLOAT
            if (left == InferredType.INT and right == InferredType.FLOAT) or (
                left == InferredType.FLOAT and right == InferredType.INT
            ):
                return InferredType.FLOAT
            # Mixed types → UNKNOWN (Python would convert, but we're strict)
            return InferredType.UNKNOWN

        # String repetition: str * int → str
        if isinstance(op, ast.Mult):
            if left == InferredType.STRING and right == InferredType.INT:
                return InferredType.STRING
            if left == InferredType.INT and right == InferredType.STRING:
                return InferredType.STRING
            if left == InferredType.INT and right == InferredType.INT:
                return InferredType.INT
            # v1.3.0: Float multiplication
            if left == InferredType.FLOAT and right == InferredType.FLOAT:
                return InferredType.FLOAT
            if (left == InferredType.INT and right == InferredType.FLOAT) or (
                left == InferredType.FLOAT and right == InferredType.INT
            ):
                return InferredType.FLOAT
            return InferredType.UNKNOWN

        # Arithmetic operators on Int → Int, Float → Float
        if isinstance(op, (ast.Sub, ast.FloorDiv, ast.Mod, ast.Pow)):
            if left == InferredType.INT and right == InferredType.INT:
                return InferredType.INT
            # v1.3.0: Float arithmetic (except FloorDiv which returns int for floats too)
            if isinstance(op, ast.FloorDiv):
                # FloorDiv returns int in Python 3 for floats too
                if left == InferredType.FLOAT or right == InferredType.FLOAT:
                    return InferredType.FLOAT  # Actually returns float in Python
            if left == InferredType.FLOAT and right == InferredType.FLOAT:
                return InferredType.FLOAT
            if (left == InferredType.INT and right == InferredType.FLOAT) or (
                left == InferredType.FLOAT and right == InferredType.INT
            ):
                return InferredType.FLOAT
            return InferredType.UNKNOWN

        # True division always returns float
        if isinstance(op, ast.Div):
            # v1.3.0: True division always returns float
            if left in (InferredType.INT, InferredType.FLOAT) and right in (
                InferredType.INT,
                InferredType.FLOAT,
            ):
                return InferredType.FLOAT
            return InferredType.UNKNOWN

        # Bitwise operators on Int → Int
        if isinstance(op, (ast.BitOr, ast.BitXor, ast.BitAnd, ast.LShift, ast.RShift)):
            if left == InferredType.INT and right == InferredType.INT:
                return InferredType.INT
            return InferredType.UNKNOWN

        # Matrix mult (@) → UNKNOWN
        if isinstance(op, ast.MatMult):
            return InferredType.UNKNOWN

        # All Python binary operators are handled above.
        # Defensive fallback for future Python operators.
        return InferredType.UNKNOWN  # pragma: no cover

    def _infer_boolop_type(self, node: ast.BoolOp) -> InferredType:
        """Infer type of boolean operation (and/or)."""
        # In Python, `and` and `or` return one of their operands
        # But for symbolic execution, we treat the result as Bool
        # (conservative but correct for our use case)

        # Check all operands are Bool
        for value in node.values:
            operand_type = self._infer_expr_type(value)
            if operand_type == InferredType.UNKNOWN:
                return InferredType.UNKNOWN

        return InferredType.BOOL
