import sys
import json
import io
import os
import types
from contextlib import redirect_stdout

class TraceEngine:
    def __init__(self, script_path):
        self.script_path = os.path.abspath(script_path)
        self.script_dir = os.path.dirname(self.script_path)
        self.trace_data = []
        self.captured_stdout = io.StringIO()
        self.step_counter = 0
        self.complexity_info = {
            "comparisons": 0,
            "swaps": 0,
            "memory_accesses": 0
        }
        
    def serialize_val(self, value):
        """
        Intelligent serializer that separates primitives from heap objects.
        Enhanced with better data structure visualization.
        """
        t = type(value).__name__
        unique_id = str(id(value))
        
        # Primitives that fit in the stack frame
        if isinstance(value, (int, float, bool, type(None))):
            return {
                "type": t,
                "value": str(value),
                "is_complex": False
            }
        
        # Strings
        if isinstance(value, str):
            if len(value) < 20:
                return {"type": t, "value": repr(value), "is_complex": False}
            else:
                return {"type": t, "value": repr(value), "id": unique_id, "is_complex": True}

        # Enhanced data structure handling
        preview = "Object"
        extra_info = {}
        
        if isinstance(value, list):
            preview = str(value)
            extra_info["length"] = len(value)
            extra_info["ds_type"] = "array"
        elif isinstance(value, tuple):
            preview = str(value)
            extra_info["length"] = len(value)
            extra_info["ds_type"] = "tuple"
        elif isinstance(value, set):
            preview = str(value)
            extra_info["length"] = len(value)
            extra_info["ds_type"] = "set"
        elif isinstance(value, dict):
            preview = str(value)
            extra_info["length"] = len(value)
            extra_info["ds_type"] = "dict"
        elif isinstance(value, (types.FunctionType, types.MethodType)):
            preview = f"function {value.__name__}(...)"
            extra_info["ds_type"] = "function"
        else:
            preview = repr(value)

        return {
            "type": t,
            "value": preview,
            "id": unique_id,
            "is_complex": True,
            **extra_info
        }

    def trace_lines(self, frame, event, arg):
        if event not in ['line', 'return', 'exception']:
            return self.trace_lines

        func_name = frame.f_code.co_name
        if func_name == '<module>': 
            func_name = 'Global'

        # Capture Locals
        variables = {}
        for k, v in frame.f_locals.items():
            if k.startswith('__'): continue
            if k == 'sys': continue
            variables[k] = self.serialize_val(v)

        # Capture Globals (selective, to show function definitions)
        globals_snapshot = {}
        if func_name != 'Global':
            for k, v in frame.f_globals.items():
                if k.startswith('__'): continue
                if isinstance(v, (types.FunctionType, types.ModuleType)):
                    globals_snapshot[k] = self.serialize_val(v)

        # Handle Return Value
        return_val = None
        if event == 'return':
            return_val = self.serialize_val(arg)

        snapshot = {
            "step": self.step_counter,
            "line": frame.f_lineno,
            "event": event,
            "func": func_name,
            "locals": variables,
            "globals": globals_snapshot,
            "return_val": return_val,
            "stdout": self.captured_stdout.getvalue(),
            "complexity": dict(self.complexity_info)
        }
        
        self.trace_data.append(snapshot)
        self.step_counter += 1
        return self.trace_lines

    def trace_calls(self, frame, event, arg):
        if event != 'call': return
        if not frame.f_code.co_filename.startswith(self.script_dir):
            return
        return self.trace_lines

    def run(self):
        with open(self.script_path, 'r') as f:
            code_str = f.read()

        code_obj = compile(code_str, self.script_path, 'exec')
        
        if self.script_dir not in sys.path:
            sys.path.insert(0, self.script_dir)

        global_scope = {"__name__": "__main__", "__file__": self.script_path}
        
        sys.settrace(self.trace_calls)
        try:
            with redirect_stdout(self.captured_stdout):
                exec(code_obj, global_scope)
        except Exception as e:
            pass 
        finally:
            sys.settrace(None)

        return code_str, self.trace_data