from __future__ import annotations 

import ast 
import builtins 
from contextlib import contextmanager 
from dataclasses import dataclass 
from typing import (
Any ,
Callable ,
Iterator ,
Literal ,
Sequence ,
TypeVar ,
cast ,
Generator ,
)

from .import argument_binding ,operator_mappings ,trace 

import itertools 
import re 


@dataclass 
class Result :
    kind :Literal ["normal","return","break","continue"]
    value :Any =None 
    is_truncated :bool =False 

    @classmethod 
    def normal (cls ,value :Any =None ,is_truncated :bool =False )->Result :
        return cls (kind ="normal",value =value ,is_truncated =is_truncated )

    @classmethod 
    def return_ (cls ,value :Any =None ,is_truncated :bool =False )->Result :
        return cls (kind ="return",value =value ,is_truncated =is_truncated )

    @classmethod 
    def break_ (cls )->Result :
        return cls (kind ="break")

    @classmethod 
    def continue_ (cls )->Result :
        return cls (kind ="continue")

    def unwrap (self )->Any :
        assert self .kind =="normal",f"Cannot unwrap {self .kind } result"
        return self .value 


@dataclass 
class FunctionDefinition :
    node :ast .FunctionDef 
    closure :dict [str ,Any ]
    interpreter :Interpreter 

    def __get__ (self ,instance :Any ,owner :type )->Any :
        if instance is None :
            return self 

        def bound_method (*args :Any ,**kwargs :Any )->Any :
            return self .interpreter ._invoke_function (
            self ,
            self .node .name ,
            [instance ]+list (args ),
            kwargs ,
            )

        return bound_method 


class IteratorWrapper :
    """
    - A wrapper class to change the output format of objects like `map` from `<map object at 0x...>` to `map<[1, 2, 3]>`
    - This class overrides the `repr` method, so it is displayed in the correct format even when output using functions like `print`
    """

    def __init__ (self ,original_type :str ,iterator :Iterator ):
        self .original_type =original_type 
        self .iterator =iterator 

    def __iter__ (self )->Iterator :
        return self .iterator 

    def __next__ (self )->Any :
        return next (self .iterator )

    def __repr__ (self )->str :
        display_iter ,_ =itertools .tee (self .iterator )
        return f"{self .original_type }<{list (display_iter )}>"


T =TypeVar ("T",bound =Callable [...,Any ])


def increment_depth (func :T )->T :
    def wrapper (self :Interpreter ,*args :Any ,**kwargs :Any )->Any :
        self .depth +=1 
        try :
            return func (self ,*args ,**kwargs )
        finally :
            self .depth -=1 

    return cast (T ,wrapper )


class Interpreter :
    def __init__ (
    self ,
    code :str ,
    max_char_length :int =128_000 *10 ,
    hide_mnemonics :bool =False ,
    override_func_defs :dict [str ,ast .FunctionDef ]|None =None ,
    ):
        self .lines =code .split ("\n")
        self .global_vars :dict [str ,Any ]={}
        self .current_char_length =0 
        self .max_char_length =max_char_length 


        self .depth =-1 
        self .trace :trace .Trace =[]
        self .is_trace_truncated =False 

        self .enable_recording =True 
        self .hide_mnemonics =hide_mnemonics 

        if override_func_defs :
            for name ,func_def_node in override_func_defs .items ():
                self .global_vars [name ]=FunctionDefinition (
                node =func_def_node ,
                closure =dict (self .global_vars ),
                interpreter =self ,
                )





    def _record (self ,mnemonic :str ,*tokens :Any )->None :
        if not self .enable_recording or self .is_trace_truncated :
            return 

        if self .hide_mnemonics :
            if not tokens :
                return 
            line =" ".join (self ._trim_obj_repr (token )for token in tokens )
        else :
            line =f"{mnemonic } "+" ".join (
            self ._trim_obj_repr (token )for token in tokens 
            )

        self .current_char_length +=len (line )
        if self .current_char_length >self .max_char_length :
            self .is_trace_truncated =True 
            return 
        self .trace .append (trace .TraceEntry (self .depth ,line ))

    @contextmanager 
    def _disable_recording (self )->Iterator [None ]:
        original =self .enable_recording 
        self .enable_recording =False 
        try :
            yield 
        finally :
            self .enable_recording =original 





    def _extract_code (self ,node :ast .AST )->str :

        assert hasattr (node ,"lineno")and isinstance (
        node .lineno ,int 
        ),"lineno is required"
        assert hasattr (node ,"col_offset"),"col_offset is required"
        line =self .lines [node .lineno -1 ]


        if isinstance (
        node ,
        (
        ast .FunctionDef ,
        ast .ClassDef ,
        ast .Try ,
        ast .With ,
        ast .For ,
        ast .If ,
        ast .Match ,
        ),
        ):
            colon_index =line .rfind (":")
            if colon_index ==-1 :
                return line [node .col_offset :]
            return line [node .col_offset :colon_index +1 ]

        end_col =(
        getattr (node ,"end_col_offset",len (line ))
        if hasattr (node ,"end_col_offset")
        else len (line )
        )
        return line [node .col_offset :end_col ]

    def _trim_obj_repr (self ,token :Any )->str :
        s =str (token )
        s =re .sub (r" at 0x[0-9a-fA-F]+","",s )
        s =re .sub (
        r"<module '([^']+)'[^>]*>",r"<module '\1'>",s 
        )
        return s 

    def _exec_statements (
    self ,
    statements :Sequence [ast .AST ],
    local_vars :dict [str ,Any ]|None ,
    )->Result :
        for statement in statements :
            self ._record ("Statement")
            result =self .visit (statement ,local_vars )
            if result .kind in ("return"  "break"  "continue"):
                return result 
        return Result .normal (is_truncated =self .is_trace_truncated )

    def _get_target_value (
    self ,
    target_id :str ,
    local_vars :dict [str ,Any ]|None ,
    )->Any :

        if local_vars and target_id in local_vars :
            return local_vars [target_id ]
        elif target_id in self .global_vars :
            return self .global_vars [target_id ]
        elif target_id in builtins .__dict__ :
            return getattr (builtins ,target_id )
        else :
            raise NameError (f"NameError: name '{target_id }' is not defined")





    def _invoke_function (
    self ,
    func :Any ,
    func_id :str ,
    positional_args :list [Any ],
    kwarg_dict :dict [str ,Any ],
    )->Any :
        if isinstance (func ,FunctionDefinition ):
            func_def_node =func .node 
            func_local_vars =argument_binding .bind_arguments (
            func_def_node ,positional_args ,kwarg_dict 
            )

            merged_locals =dict (func .closure )
            merged_locals .update (func_local_vars )

            call_result =self ._exec_statements (func_def_node .body ,merged_locals )
            if call_result .kind =="return":
                return call_result .value 
            elif call_result .kind =="normal":
                return None 
            else :
                raise SyntaxError (f"{call_result .kind } statement outside of loop")
        else :
            if func is None :
                raise NameError (f"NameError: name '{func_id }' is not defined")
            return func (*positional_args ,**kwarg_dict )





    def _bind_comprehension_target (
    self ,
    target_node :ast .AST ,
    value :Any ,
    local_vars :dict [str ,Any ],
    )->None :
        match target_node :
            case ast .Name (id =var_name ):
                local_vars [var_name ]=value 
            case ast .Tuple (elts =elts ):
                if not isinstance (value ,(list ,tuple )):
                    try :
                        rhs_value =list (value )
                    except TypeError :
                        raise TypeError (f"Cannot unpack non-sequence: {rhs_value !r }")
                if len (elts )!=len (value ):
                    raise ValueError ("Unpack size does not match")
                for e ,v in zip (elts ,value ):
                    self ._bind_comprehension_target (e ,v ,local_vars )
            case _ :
                raise NotImplementedError (
                f"Unsupported comprehension target: {target_node .__class__ .__name__ }"
                )

    def _evaluate_listcomp (
    self ,
    elt :ast .expr ,
    generators :list [ast .comprehension ],
    local_vars :dict [str ,Any ]|None ,
    )->list [Any ]:
        result =[]

        def visit_comprehension (gen_index :int ,local_env :dict [str ,Any ])->None :
            """Recursively handle each generator in the list comprehension."""
            if gen_index >=len (generators ):

                val =self .visit (elt ,local_env ).unwrap ()
                result .append (val )
                return 

            gen =generators [gen_index ]
            iter_obj =self .visit (gen .iter ,local_env ).unwrap ()
            try :
                iterator =builtins .iter (iter_obj )
            except TypeError :
                raise TypeError (f"Object {iter_obj } is not iterable")
            for item in iterator :
                loop_env =dict (local_env )
                self ._bind_comprehension_target (gen .target ,item ,loop_env )
                should_continue =True 
                for if_node in gen .ifs :
                    cond_val =self .visit (if_node ,loop_env ).unwrap ()
                    if not cond_val :
                        should_continue =False 
                        break 
                if should_continue :
                    visit_comprehension (gen_index +1 ,loop_env )

        visit_comprehension (0 ,local_vars or {})
        return result 

    def _evaluate_dictcomp (
    self ,
    key_expr :ast .expr ,
    value_expr :ast .expr ,
    generators :list [ast .comprehension ],
    local_env :dict [str ,Any ]|None ,
    )->dict [Any ,Any ]:
        result ={}

        def visit_comp (gen_index :int ,env :dict [str ,Any ])->None :
            if gen_index >=len (generators ):
                key =self .visit (key_expr ,env ).unwrap ()
                value =self .visit (value_expr ,env ).unwrap ()
                result [key ]=value 
                return 
            gen =generators [gen_index ]
            iter_val =self .visit (gen .iter ,env ).unwrap ()
            for item in iter_val :
                loop_env =dict (env )
                self ._bind_comprehension_target (gen .target ,item ,loop_env )
                passed_all =True 
                for if_node in gen .ifs :
                    if_cond =self .visit (if_node ,loop_env ).unwrap ()
                    if not if_cond :
                        passed_all =False 
                        break 
                if passed_all :
                    visit_comp (gen_index +1 ,loop_env )

        visit_comp (0 ,local_env or {})
        return result 

    def _evaluate_generatorexp (
    self ,
    elt :ast .expr ,
    generators :list [ast .comprehension ],
    local_vars :dict [str ,Any ]|None ,
    )->Generator [Any ,None ,None ]:
        """
        Build a generator (iterator) that lazily evaluates the expression 'elt'
        for each combination of generator loops and if-conditions.
        """


        def gen ()->Generator [Any ,None ,None ]:
            def visit_generator_level (
            gen_index :int ,env :dict [str ,Any ]
            )->Generator [Any ,None ,None ]:
                if gen_index >=len (generators ):

                    val =self .visit (elt ,env ).unwrap ()
                    yield val 
                    return 
                gen_node =generators [gen_index ]
                iter_obj =self .visit (gen_node .iter ,env ).unwrap ()
                for item in iter_obj :
                    loop_env =dict (env )
                    self ._bind_comprehension_target (gen_node .target ,item ,loop_env )

                    passed_all =True 
                    for if_node in gen_node .ifs :
                        cond_val =self .visit (if_node ,loop_env ).unwrap ()
                        if not cond_val :
                            passed_all =False 
                            break 
                    if passed_all :
                        yield from visit_generator_level (gen_index +1 ,loop_env )

            yield from visit_generator_level (0 ,local_vars or {})

        return gen ()





    def _handle_try (
    self ,
    body :Sequence [ast .AST ],
    handlers :Sequence [ast .ExceptHandler ],
    orelse :Sequence [ast .AST ],
    finalbody :Sequence [ast .AST ],
    local_vars :dict [str ,Any ]|None ,
    )->Result :
        exception_raised :BaseException |None =None 
        handled =False 
        try :
            result =self ._exec_statements (body ,local_vars )
            if result .kind in ("return","break","continue"):
                final_result =self ._exec_statements (finalbody ,local_vars )
                if final_result .kind in ("return","break","continue"):
                    return final_result 
                return result 
        except Exception as e :
            exception_raised =e 

        if exception_raised is not None :
            for handler in handlers :
                if handler .type is None or self ._exception_matches (
                handler .type ,exception_raised ,local_vars 
                ):
                    if handler .name :
                        target_scope =(
                        local_vars if local_vars is not None else self .global_vars 
                        )
                        target_scope [handler .name ]=exception_raised 
                    try :
                        result =self ._exec_statements (handler .body ,local_vars )
                        if result .kind in ("return","break","continue"):
                            final_result =self ._exec_statements (finalbody ,local_vars )
                            if final_result .kind in ("return","break","continue"):
                                return final_result 
                            return result 
                        handled =True 
                        break 
                    except Exception as e2 :
                        exception_raised =e2 
                        handled =False 
                        break 

        if exception_raised is None and not handled :
            result =self ._exec_statements (orelse ,local_vars )
            if result .kind in ("return","break","continue"):

                final_result =self ._exec_statements (finalbody ,local_vars )
                if final_result .kind in ("return","break","continue"):
                    return final_result 
                return result 


        final_result =self ._exec_statements (finalbody ,local_vars )
        if final_result .kind in ("return","break","continue"):
            return final_result 

        if exception_raised is not None and not handled :
            raise exception_raised 

        return Result .normal ()

    def _exception_matches (
    self ,
    type_node :ast .AST ,
    exception :BaseException ,
    local_vars :dict [str ,Any ]|None ,
    )->bool :
        exception_type =self .visit (type_node ,local_vars ).unwrap ()
        return isinstance (exception ,exception_type )




    @increment_depth 
    def visit (self ,node :ast .AST ,local_vars :dict [str ,Any ]|None )->Result :
        match node :



            case ast .Module (body =body ):
                assert local_vars is None ,"Module should not have local_vars"


                return self ._exec_statements (node .body ,None )

            case ast .FunctionDef ():
                self ._record ("FunctionDef",self ._extract_code (node ).strip ())

                func_def_obj :FunctionDefinition =FunctionDefinition (
                node =node ,
                closure =local_vars if local_vars is not None else {},
                interpreter =self ,
                )

                decorated =func_def_obj 
                for decorated_node in reversed (
                node .decorator_list 
                ):
                    decorated_callable =self .visit (decorated_node ,local_vars ).unwrap ()
                    decorated =self ._invoke_function (
                    decorated_callable ,"<decorator>",[decorated ],{}
                    )
                    self ._record ("Decorator",f"{decorated_callable } -> {decorated }")

                target_scope =(
                local_vars if local_vars is not None else self .global_vars 
                )
                target_scope [node .name ]=decorated 
                return Result .normal ()

            case ast .Global (names =names ):
                self ._record ("Global",self ._extract_code (node ).strip ())
                if local_vars is not None :
                    forced_globals =local_vars .setdefault ("__forced_globals__",set ())
                    for name in names :
                        forced_globals .add (name )
                return Result .normal ()




            case ast .Name ():
                value =self ._get_target_value (node .id ,local_vars )
                self ._record ("Name",f"{node .id } = {value }")
                return Result .normal (value )

            case ast .Assign (targets =targets ,value =value ):
                self ._record ("Assign",self ._extract_code (node ).strip ())
                forced_globals =(
                local_vars .get ("__forced_globals__",set ())if local_vars else set ()
                )

                rhs_value =self .visit (value ,local_vars ).unwrap ()


                for lhs in targets :
                    match lhs :
                        case ast .Name (id =id ):
                            target_scope =(
                            self .global_vars 
                            if (local_vars is None or id in forced_globals )
                            else local_vars 
                            )
                            target_scope [id ]=rhs_value 
                            self ._record ("Assign",f"{id } = {rhs_value }")
                        case ast .Subscript (value =sub_val ,slice =sub_slice ):
                            with self ._disable_recording ():
                                container =self .visit (sub_val ,local_vars ).unwrap ()
                                index =self .visit (sub_slice ,local_vars ).unwrap ()
                            container [index ]=rhs_value 
                            if isinstance (sub_val ,ast .Name ):
                                container_name =sub_val .id 
                            else :
                                container_name =self ._extract_code (sub_val ).strip ()

                            self ._record (
                            "Assign",f"{container_name }[{index }] = {rhs_value }"
                            )
                        case ast .Tuple (elts =elts ):

                            if not isinstance (rhs_value ,(list ,tuple )):
                                try :
                                    rhs_value =list (rhs_value )
                                except TypeError :
                                    raise TypeError (
                                    f"Cannot unpack non-sequence: {rhs_value !r }"
                                    )
                            if len (elts )!=len (rhs_value ):
                                raise ValueError (
                                f"Unpack mismatch: {len (elts )} variables but got {len (rhs_value )} values"
                                )
                            for sub_lhs ,sub_val in zip (elts ,rhs_value ):

                                match sub_lhs :
                                    case ast .Name (id =sub_id ):
                                        sub_scope =(
                                        self .global_vars 
                                        if (
                                        local_vars is None 
                                        or sub_id in forced_globals 
                                        )
                                        else local_vars 
                                        )
                                        sub_scope [sub_id ]=sub_val 
                                        self ._record ("Assign",f"{sub_id } = {sub_val }")
                                    case ast .Subscript (
                                    value =sub_val_node ,slice =sub_slice_node 
                                    ):
                                        with self ._disable_recording ():
                                            container2 =self .visit (
                                            sub_val_node ,local_vars 
                                            ).unwrap ()
                                            index2 =self .visit (
                                            sub_slice_node ,local_vars 
                                            ).unwrap ()
                                        container2 [index2 ]=sub_val 
                                        self ._record (
                                        "AssignSubscript",
                                        f"{self ._extract_code (sub_lhs ).strip ()} = {sub_val }",
                                        )
                                    case _ :
                                        raise NotImplementedError (
                                        f"Assignment to nested {sub_lhs .__class__ .__name__ } not implemented."
                                        )
                        case ast .Attribute (value =obj ,attr =attr_name ):
                            obj_value =self .visit (obj ,local_vars ).unwrap ()
                            setattr (obj_value ,attr_name ,rhs_value )
                            self ._record (
                            "Assign",f"{obj_value }.{attr_name } = {rhs_value }"
                            )
                        case _ :
                            raise NotImplementedError (
                            f"Assignment to {lhs .__class__ .__name__ } not implemented."
                            )
                return Result .normal ()

            case ast .AnnAssign (target =target ,value =value ):
                self ._record ("AnnAssign",self ._extract_code (node ).strip ())
                assigned_value =(
                self .visit (value ,local_vars ).unwrap ()
                if value is not None 
                else None 
                )

                match target :
                    case ast .Name (id =var_name ):
                        forced_globals =(
                        local_vars .get ("__forced_globals__",set ())
                        if local_vars 
                        else set ()
                        )
                        target_scope =(
                        self .global_vars 
                        if (local_vars is None or var_name in forced_globals )
                        else local_vars 
                        )
                        target_scope [var_name ]=assigned_value 
                        self ._record ("AnnAssign",f"{var_name } = {assigned_value }")
                    case _ :
                        raise NotImplementedError (
                        f"Unsupported AnnAssign target: {target .__class__ .__name__ }"
                        )
                return Result .normal ()

            case ast .AugAssign (target =target ,op =op ,value =value ):
                self ._record ("AugAssign",self ._extract_code (node ).strip ())

                match target :
                    case ast .Name (id =name ):
                        old_value =self ._get_target_value (name ,local_vars )

                        right_value =self .visit (value ,local_vars ).unwrap ()

                        if isinstance (op ,ast .Add )and isinstance (
                        old_value ,list 
                        ):
                            try :
                                new_items =list (right_value )
                            except TypeError :
                                raise TypeError (
                                f"Can only extend list with an iterable (not '{type (right_value ).__name__ }')"
                                )
                            old_value .extend (new_items )
                            new_value =old_value 
                        else :
                            op_type =type (op )
                            try :
                                bin_op ,op_str =operator_mappings ._BIN_OPS [op_type ]
                                new_value =bin_op (old_value ,right_value )
                            except KeyError :
                                raise NotImplementedError (
                                f"Not implemented: {op_type .__name__ }"
                                )


                        if local_vars is not None and name in local_vars :
                            local_vars [name ]=new_value 
                        else :
                            self .global_vars [name ]=new_value 

                        self ._record ("AugAssign",f"{name } = {new_value }")


                    case ast .Subscript (value =sub_val ,slice =sub_slice ):
                        self ._record (
                        "AugAssignSubscript",self ._extract_code (target ).strip ()
                        )

                        with self ._disable_recording ():
                            container =self .visit (sub_val ,local_vars ).unwrap ()
                            index =self .visit (sub_slice ,local_vars ).unwrap ()

                        old_value =container [index ]
                        right_value =self .visit (value ,local_vars ).unwrap ()

                        if isinstance (old_value ,list )and isinstance (
                        op ,ast .Add 
                        ):
                            try :
                                new_items =list (right_value )
                            except TypeError :
                                raise TypeError (
                                f"Can only extend list with an iterable (not '{type (right_value ).__name__ }')"
                                )
                            old_value .extend (new_items )
                            new_value =old_value 
                        else :
                            op_type =type (op )
                            try :
                                bin_op ,op_str =operator_mappings ._BIN_OPS [op_type ]
                                new_value =bin_op (old_value ,right_value )
                            except KeyError :
                                raise NotImplementedError (
                                f"Not implemented: {op_type .__name__ }"
                                )

                        container [index ]=new_value 
                        self ._record (
                        "AugAssign",
                        f"{self ._extract_code (target ).strip ()} = {new_value }",
                        )

                    case _ :
                        raise NotImplementedError (
                        f"AugAssign target not supported: {target .__class__ .__name__ }"
                        )

                return Result .normal ()

            case ast .Delete (targets =targets ):
                self ._record ("Delete",self ._extract_code (node ).strip ())

                for t in targets :
                    match t :
                        case ast .Name (id =name ):

                            if local_vars is not None and name in local_vars :
                                del local_vars [name ]
                            elif name in self .global_vars :
                                del self .global_vars [name ]
                            else :
                                raise NameError (
                                f"NamError: name '{name }' is not defined"
                                )
                            self ._record ("DeleteName",name )

                        case ast .Subscript (value =sub_val ,slice =sub_slice ):
                            with self ._disable_recording ():
                                container =self .visit (sub_val ,local_vars ).unwrap ()
                                index =self .visit (sub_slice ,local_vars ).unwrap ()
                            del container [index ]
                            self ._record ("DeleteSubscript",f"{container }[{index }]")

                        case ast .Attribute (value =attr_val ,attr =attr_name ):
                            obj =self .visit (attr_val ,local_vars ).unwrap ()
                            delattr (obj ,attr_name )
                            self ._record ("DeleteAttribute",f"{obj }.{attr_name }")

                        case _ :
                            raise NotImplementedError (
                            f"Delete target not supported: {t .__class__ .__name__ }"
                            )

                return Result .normal ()




            case ast .Expr (value =value ):
                self ._record ("Expr",self ._extract_code (node ).strip ())
                return self .visit (value ,local_vars )

            case ast .UnaryOp ():
                self ._record ("UnaryOp",self ._extract_code (node ).strip ())
                operand =self .visit (node .operand ,local_vars ).unwrap ()
                match node .op :
                    case ast .USub ():
                        op_str ="-"
                        value =-operand 
                    case ast .Not ():
                        op_str ="Not"
                        value =not operand 
                    case ast .Invert ():
                        op_str ="~"
                        value =~operand 
                    case _ :
                        raise NotImplementedError (
                        f"Not implemented: {node .op .__class__ .__name__ }"
                        )
                return Result .normal (value )

            case ast .BinOp (left =left_expr ,op =op ,right =right_expr ):
                self ._record ("BinOp",self ._extract_code (node ).strip ())
                lhs =self .visit (left_expr ,local_vars ).unwrap ()
                self ._record ("BinOpLeft",lhs )
                rhs =self .visit (right_expr ,local_vars ).unwrap ()
                self ._record ("BinOpRight",rhs )

                op_type =type (op )
                try :
                    bin_op ,op_str =operator_mappings ._BIN_OPS [op_type ]
                    value =bin_op (lhs ,rhs )
                except KeyError :
                    raise NotImplementedError (f"Not implemented: {op_type .__name__ }")
                self ._record ("BinOp",f"{lhs } {op_str } {rhs } = {value }")
                return Result .normal (value )

            case ast .BoolOp (op =op ,values =values ):
                self ._record ("BoolOp",self ._extract_code (node ).strip ())

                match op :
                    case ast .And ():
                        current_val =self .visit (values [0 ],local_vars ).unwrap ()
                        for val_node in values [1 :]:
                            if not current_val :

                                break 
                            current_val =self .visit (val_node ,local_vars ).unwrap ()
                        return Result .normal (current_val )

                    case ast .Or ():
                        current_val =self .visit (values [0 ],local_vars ).unwrap ()
                        for val_node in values [1 :]:
                            if current_val :

                                break 
                            current_val =self .visit (val_node ,local_vars ).unwrap ()
                        return Result .normal (current_val )

                    case _ :
                        raise NotImplementedError (
                        f"Unsupported BoolOp: {op .__class__ .__name__ }"
                        )

            case ast .Compare (left =left ,ops =cmpops ,comparators =comparators ):
                self ._record ("Compare",self ._extract_code (node ).strip ())

                self ._record ("CompareLeft",self ._extract_code (left ).strip ())
                lhs_value =self .visit (left ,local_vars ).unwrap ()

                final_result =True 
                current_lhs =lhs_value 

                for cmpop ,cmp_node in zip (cmpops ,comparators ):
                    self ._record ("CompareRight",self ._extract_code (cmp_node ).strip ())
                    rhs_value =self .visit (cmp_node ,local_vars ).unwrap ()

                    cmp_op_type =type (cmpop )
                    try :
                        cmp_op ,op_str =operator_mappings ._CMP_OPS [cmp_op_type ]
                        cmp_result =cmp_op (current_lhs ,rhs_value )
                    except KeyError :
                        raise NotImplementedError (
                        f"Not implemented: {cmp_op_type .__name__ }"
                        )

                    self ._record (
                    "CompareResult",
                    f"{current_lhs } {op_str } {rhs_value } = {cmp_result }",
                    )

                    if not cmp_result :
                        final_result =False 
                        break 

                    current_lhs =rhs_value 

                return Result .normal (final_result )

            case ast .Call (func =func ,args =args ,keywords =keywords ):
                self ._record ("Call",self ._extract_code (node ).strip ())

                if isinstance (func ,ast .Name ):
                    func_obj =None 
                    func_id =func .id 
                elif isinstance (func ,ast .Attribute ):
                    func_obj =self .visit (func .value ,local_vars ).unwrap ()
                    func_id =func .attr 
                else :
                    raise NotImplementedError (
                    f"Call target not supported: {type (func )}"
                    )


                positional_args =[]
                for i ,arg_node in enumerate (args ):
                    self ._record (f"CallArg{i }",self ._extract_code (arg_node ).strip ())
                    if isinstance (
                    arg_node ,ast .Starred 
                    ):
                        starred_items =self .visit (arg_node .value ,local_vars ).unwrap ()
                        try :
                            iter (starred_items )
                        except TypeError :
                            raise TypeError (
                            f"Cannot use * on a non-iterable object: {type (starred_items ).__name__ }"
                            )

                        positional_args .extend (starred_items )
                    else :
                        positional_args .append (
                        self .visit (arg_node ,local_vars ).unwrap ()
                        )


                kwarg_dict ={}
                for kw_node in keywords :
                    if kw_node .arg is None :
                        unpack_dict =self .visit (kw_node .value ,local_vars ).unwrap ()
                        if not isinstance (unpack_dict ,dict ):
                            raise TypeError (
                            f"Keyword argument unpacking expects a dict, got {type (unpack_dict )}"
                            )
                        kwarg_dict .update (unpack_dict )
                    else :
                        kwarg_name =kw_node .arg 
                        kwarg_value =self .visit (kw_node .value ,local_vars ).unwrap ()
                        kwarg_dict [kwarg_name ]=kwarg_value 

                possible_func =(
                self ._get_target_value (func_id ,local_vars )
                if func_obj is None 
                else getattr (func_obj ,func_id )
                )
                value =self ._invoke_function (
                possible_func ,func_id ,positional_args ,kwarg_dict 
                )

                if isinstance (value ,Iterator ):
                    original_type =type (value ).__name__ 
                    _ ,value_copy =itertools .tee (value )
                    value =IteratorWrapper (original_type ,value_copy )

                self ._record (
                "Call",
                f"{func_id }({', '.join (map (repr ,positional_args ))}, **{kwarg_dict }) = {value }"
                if kwarg_dict 
                else f"{func_id }({', '.join (map (repr ,positional_args ))}) = {value }",
                )

                return Result .normal (value )

            case ast .Attribute (value =attr_value ,attr =attr_name ):
                self ._record ("Attribute",self ._extract_code (node ).strip ())
                obj =self .visit (attr_value ,local_vars ).unwrap ()
                try :
                    value =getattr (obj ,attr_name )
                except AttributeError :
                    raise AttributeError (
                    f"Object '{obj }' has no attribute '{attr_name }'"
                    )
                self ._record ("Attribute",f"{obj }.{attr_name } = {value }")
                return Result .normal (value )

            case ast .Subscript (value =value ,slice =slice ):
                self ._record ("Subscript",self ._extract_code (node ).strip ())
                container_name =self ._extract_code (value ).strip ()
                with self ._disable_recording ():
                    container =self .visit (value ,local_vars ).unwrap ()
                index =self .visit (slice ,local_vars ).unwrap ()
                self ._record ("Subscript",f"{container_name }[{index }]")
                return Result .normal (container [index ])

            case ast .Slice (lower =lower ,upper =upper ,step =step ):
                self ._record ("Slice",self ._extract_code (node ).strip ())
                if lower is not None :
                    self ._record ("SliceLower",self ._extract_code (lower ).strip ())
                    lower_val =self .visit (lower ,local_vars ).unwrap ()
                else :
                    lower_val =None 

                if upper is not None :
                    self ._record ("SliceUpper",self ._extract_code (upper ).strip ())
                    upper_val =self .visit (upper ,local_vars ).unwrap ()
                else :
                    upper_val =None 

                if step is not None :
                    self ._record ("SliceStep",self ._extract_code (step ).strip ())
                    step_val =self .visit (step ,local_vars ).unwrap ()
                else :
                    step_val =None 

                result_slice =builtins .slice (lower_val ,upper_val ,step_val )
                self ._record ("Slice",f"-> {result_slice }")
                return Result .normal (result_slice )

            case ast .NamedExpr (target =target ,value =value_expr ):
                self ._record ("NamedExpr",self ._extract_code (node ).strip ())
                val =self .visit (value_expr ,local_vars ).unwrap ()
                match target :
                    case ast .Name (id =name ):
                        if local_vars is not None and name in local_vars :
                            local_vars [name ]=val 
                        else :
                            self .global_vars [name ]=val 
                    case _ :
                        raise NotImplementedError (
                        f"Unsupported NamedExpr target: {target .__class__ .__name__ }"
                        )
                return Result .normal (val )

            case ast .Lambda (args =lambda_args ,body =lambda_body ):
                self ._record ("Lambda",self ._extract_code (node ).strip ())


                param_names =[arg .arg for arg in lambda_args .args ]

                def lambda_function (*call_args :tuple )->Any :
                    lambda_local_vars ={
                    name :val for name ,val in zip (param_names ,call_args )
                    }
                    merged_locals =local_vars or {}
                    merged_locals .update (lambda_local_vars )

                    result =self .visit (lambda_body ,merged_locals )
                    return result .unwrap ()

                return Result .normal (lambda_function )




            case ast .Constant (value =value ):
                self ._record ("Constant",repr (value ))
                return Result .normal (value )

            case ast .List (elts =elts ):
                self ._record ("List",self ._extract_code (node ).strip ())
                rets =[]
                for i ,elt in enumerate (elts ):
                    self ._record (f"List{i }",self ._extract_code (elt ).strip ())
                    rets .append (self .visit (elt ,local_vars ).unwrap ())
                return Result .normal (rets )

            case ast .Set (elts =elts ):
                self ._record ("Set",self ._extract_code (node ).strip ())
                result_set =set ()
                for i ,elt in enumerate (elts ):
                    val =self .visit (elt ,local_vars ).unwrap ()
                    result_set .add (val )
                return Result .normal (result_set )

            case ast .Tuple (elts =elements ):
                self ._record ("Tuple",self ._extract_code (node ).strip ())
                items =[]
                for i ,elt in enumerate (elements ):
                    value =self .visit (elt ,local_vars ).unwrap ()
                    items .append (value )
                result_tuple =tuple (items )
                return Result .normal (result_tuple )

            case ast .Dict (keys =dict_keys ,values =dict_values ):
                self ._record ("Dict",self ._extract_code (node ).strip ())
                d ={}
                for key_node ,value_node in zip (dict_keys ,dict_values ):
                    if key_node is None :
                        unpack_dict =self .visit (value_node ,local_vars ).unwrap ()
                        if not isinstance (unpack_dict ,dict ):
                            raise TypeError (
                            f"Dict unpacking expects a dict, got {type (unpack_dict )}"
                            )
                        d .update (unpack_dict )
                    else :
                        key_val =self .visit (key_node ,local_vars ).unwrap ()
                        value_val =self .visit (value_node ,local_vars ).unwrap ()
                        d [key_val ]=value_val 
                return Result .normal (d )

            case ast .ListComp (elt =elt ,generators =generators ):
                self ._record ("ListComp",self ._extract_code (node ).strip ())
                result_list =self ._evaluate_listcomp (elt ,generators ,local_vars )
                self ._record ("ListCompResult",result_list )
                return Result .normal (result_list )

            case ast .DictComp (key =key_expr ,value =value_expr ,generators =generators ):
                self ._record ("DictComp",self ._extract_code (node ).strip ())
                result_dict =self ._evaluate_dictcomp (
                key_expr ,value_expr ,generators ,local_vars 
                )
                self ._record ("DictCompResult",result_dict )
                return Result .normal (result_dict )

            case ast .GeneratorExp (elt =elt ,generators =generators ):
                self ._record ("GeneratorExp",self ._extract_code (node ).strip ())
                result_iter =self ._evaluate_generatorexp (elt ,generators ,local_vars )
                return Result .normal (result_iter )

            case ast .JoinedStr (values =string_parts ):
                self ._record ("JoinedStr",self ._extract_code (node ).strip ())
                result_parts =[]
                for part in string_parts :
                    match part :
                        case ast .Constant (value =cval ):
                            if not isinstance (cval ,str ):
                                raise TypeError (
                                f"JoinedStr has non-string Constant: {cval }"
                                )
                            result_parts .append (cval )

                        case ast .FormattedValue (value =val_node ,conversion =conversion ):
                            value =self .visit (val_node ,local_vars ).unwrap ()

                            if conversion ==-1 :
                                val_str =str (value )
                            elif conversion ==115 :
                                val_str =str (value )
                            elif conversion ==114 :
                                val_str =repr (value )
                            else :
                                raise NotImplementedError (
                                f"Unsupported conversion code: {conversion }"
                                )
                            result_parts .append (val_str )

                        case _ :
                            raise NotImplementedError (
                            f"Unexpected node in JoinedStr: {part }"
                            )
                joined_str ="".join (result_parts )
                self ._record ("JoinedStr",f"{joined_str !r }")
                return Result .normal (joined_str )




            case ast .If (test =test ,body =body ,orelse =orelse ):
                self ._record ("If",self ._extract_code (node ).strip ())
                cond_val =self .visit (test ,local_vars ).unwrap ()
                self ._record ("IfCond",cond_val )
                if_body_result =self ._exec_statements (
                body if cond_val else orelse ,local_vars 
                )
                if if_body_result .kind in ("return","break","continue"):
                    return if_body_result 
                return Result .normal ()

            case ast .IfExp (test =test ,body =body ,orelse =orelse ):
                self ._record ("IfExp",self ._extract_code (node ).strip ())
                cond_val =self .visit (test ,local_vars ).unwrap ()
                self ._record ("IfExpCond",cond_val )
                return Result .normal (
                self .visit (body if cond_val else orelse ,local_vars ).unwrap ()
                )

            case ast .While (test =test ,body =body ):
                self ._record ("While",self ._extract_code (node ).strip ())
                while True :
                    self ._record ("WhileCond",self ._extract_code (test ).strip ())
                    test_val =self .visit (test ,local_vars ).unwrap ()
                    self ._record ("WhileCondResult",test_val )
                    if not test_val :
                        break 
                    loop_broken =False 
                    for statement in body :
                        self ._record ("WhileBody",self ._extract_code (statement ).strip ())
                        while_body_result =self .visit (statement ,local_vars )
                        if while_body_result .kind =="return":
                            return while_body_result 
                        if while_body_result .kind =="break":
                            loop_broken =True 
                            break 
                        if while_body_result .kind =="continue":
                            break 
                    if loop_broken :
                        break 
                return Result .normal ()

            case ast .For (target =target ,iter =iter_ ,body =body ,orelse =orelse ):
                self ._record ("For",self ._extract_code (node ).strip ())
                iterable_obj =self .visit (iter_ ,local_vars ).unwrap ()
                try :
                    iterator =iter (iterable_obj )
                except TypeError :
                    raise TypeError (f"Object '{iterable_obj }' is not iterable")
                finished_normally =True 
                for item in iterator :
                    self ._bind_comprehension_target (
                    target ,item ,local_vars or self .global_vars 
                    )
                    for_body_result =self ._exec_statements (body ,local_vars )
                    if for_body_result .kind =="return":
                        return for_body_result 
                    elif for_body_result .kind =="break":
                        finished_normally =False 
                        break 
                    elif for_body_result .kind =="continue":
                        continue 
                if finished_normally :
                    for_body_result =self ._exec_statements (orelse ,local_vars )
                    if for_body_result .kind in ("return","break"):
                        return for_body_result 
                return Result .normal ()

            case ast .Break ():
                self ._record ("Break",self ._extract_code (node ).strip ())
                return Result .break_ ()

            case ast .Continue ():
                self ._record ("Continue",self ._extract_code (node ).strip ())
                return Result .continue_ ()

            case ast .Return (value =value ):
                self ._record ("Return",self ._extract_code (node ).strip ())
                if value is None :
                    return Result .return_ ()
                return Result .return_ (
                self .visit (value ,local_vars ).unwrap (),
                is_truncated =self .is_trace_truncated ,
                )

            case ast .Assert (test =test_expr ,msg =msg_expr ):
                self ._record ("Assert",self ._extract_code (node ).strip ())

                cond =self .visit (test_expr ,local_vars ).unwrap ()
                if not cond :
                    msg_val =None 
                    if msg_expr is not None :
                        msg_val =self .visit (msg_expr ,local_vars ).unwrap ()
                    raise AssertionError (msg_val )

                return Result .normal ()

            case ast .Try (
            body =body ,handlers =handlers ,orelse =orelse ,finalbody =finalbody 
            ):
                self ._record ("Try",self ._extract_code (node ).strip ())
                return self ._handle_try (body ,handlers ,orelse ,finalbody ,local_vars )

            case ast .Pass ():
                self ._record ("Pass",self ._extract_code (node ).strip ())
                return Result .normal ()




            case ast .ClassDef (
            name =name ,
            bases =bases ,
            keywords =keywords ,
            body =body ,
            decorator_list =decorators ,
            ):
                self ._record ("ClassDef",self ._extract_code (node ).strip ())


                base_classes =(
                [self .visit (base ,local_vars ).unwrap ()for base in bases ]
                if bases 
                else [object ]
                )

                class_namespace :dict [str ,Any ]={}
                class_namespace ["__module__"]="__main__"


                self ._exec_statements (body ,class_namespace )



                new_class =type (name ,tuple (base_classes ),class_namespace )

                for decorated_node in reversed (decorators ):
                    decorated_callable =self .visit (decorated_node ,local_vars ).unwrap ()
                    new_class =self ._invoke_function (
                    decorated_callable ,"<class-decorator>",[new_class ],{}
                    )
                    self ._record ("Decorator",f"{decorated_callable } -> {new_class }")

                target_scope =(
                local_vars if local_vars is not None else self .global_vars 
                )
                target_scope [name ]=new_class 

                return Result .normal (new_class )




            case ast .Import (names =names ):
                self ._record ("Import",self ._extract_code (node ).strip ())

                for alias in names :

                    module_name =alias .name 
                    local_name =alias .asname if alias .asname else module_name 


                    try :
                        imported_module =__import__ (module_name )
                    except ImportError :
                        raise ImportError (f"Cannot import module '{module_name }'")


                    self .global_vars [local_name ]=imported_module 

                    self ._record ("Import",f"import {module_name } as {local_name }")

                return Result .normal ()

            case ast .ImportFrom (module =module_name ,names =names ,level =level ):
                self ._record ("ImportFrom",self ._extract_code (node ).strip ())

                if level !=0 :
                    raise NotImplementedError ("Relative imports not implemented yet.")


                try :
                    imported_module =__import__ (module_name )
                except ImportError :
                    raise ImportError (f"Cannot import module '{module_name }'")

                for alias in names :
                    public_name =alias .name 
                    local_name =alias .asname if alias .asname else public_name 
                    if public_name =="*":
                        __all__ =getattr (imported_module ,"__all__",None )
                        if __all__ is None :

                            all_names =[
                            n for n in dir (imported_module )if not n .startswith ("_")
                            ]
                        else :
                            all_names =__all__ 
                        for name in all_names :
                            self .global_vars [name ]=getattr (imported_module ,name )
                            self ._record (
                            "ImportFrom*",f"from {module_name } import {name }"
                            )
                    else :
                        try :
                            imported_obj =getattr (imported_module ,public_name )
                        except AttributeError :
                            raise ImportError (
                            f"Cannot import name '{public_name }' from '{module_name }'"
                            )

                        self .global_vars [local_name ]=imported_obj 
                        self ._record (
                        "ImportFrom",
                        f"from {module_name } import {public_name } as {local_name }",
                        )
                return Result .normal ()

            case _ :
                self ._record ("NotImplemented",node )
                raise NotImplementedError (f"Not implemented: {node .__class__ .__name__ }")

        assert False ,"unreachable"
