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

class TraceEngine:
    def __init__(self, script_path):
        # Normalize path for consistent comparison
        self.script_path = os.path.abspath(script_path)
        self.trace_data = []
        self.captured_stdout = io.StringIO()
        self.step_counter = 0
        
    def safe_repr(self, obj):
        try:
            if isinstance(obj, (int, float, str, bool, type(None))):
                return obj
            if isinstance(obj, (list, tuple, set)):
                return str(obj)
            if isinstance(obj, dict):
                return str(obj)
            return repr(obj)
        except Exception:
            return "<Unserializable Object>"

    def get_locals(self, frame):
        variables = {}
        for key, value in frame.f_locals.items():
            if key.startswith('__'): continue
            if key == 'sys': continue
            variables[key] = {
                "type": type(value).__name__,
                "value": self.safe_repr(value)
            }
        return variables

    def trace_calls(self, frame, event, arg):
        if event != 'call':
            return
        
        co = frame.f_code
        filename = co.co_filename
        
        # RELAXED CHECK: Use endswith to be safer against path variations
        if not filename.endswith(os.path.basename(self.script_path)):
            return
            
        return self.trace_lines

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

        # Capture stdout up to this point
        current_stdout = self.captured_stdout.getvalue()

        snapshot = {
            "step": self.step_counter,
            "line": frame.f_lineno,
            "event": event,
            "locals": self.get_locals(frame),
            "stdout": current_stdout
        }
        
        self.trace_data.append(snapshot)
        self.step_counter += 1
        return self.trace_lines

    def run(self):
        """Executes the script within the trace context."""
        with open(self.script_path, 'r') as f:
            code_str = f.read()

        # --- THE CRITICAL FIX ---
        # We must compile the code with the filename explicitly 
        # so sys.settrace sees "test.py" instead of "<string>"
        code_obj = compile(code_str, self.script_path, 'exec')
        # ------------------------

        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:
            self.trace_data.append({
                "step": self.step_counter,
                "line": 0,
                "event": "exception",
                "locals": {},
                "stdout": self.captured_stdout.getvalue() + f"\nError: {str(e)}"
            })
        finally:
            sys.settrace(None)

        return code_str, self.trace_data