import time
from typing import Optional, List, Dict, Any
from dataclasses import dataclass, field


@dataclass
class ExecutionTrace:
    segment_id: str
    byte_start: int
    byte_end: int
    duration: float
    exception: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


class ExecutionTracer:
    def __init__(self, verbose: bool = False):
        self.verbose = verbose
        self.traces: List[ExecutionTrace] = []

    def trace_segment(self, segment_id: str, byte_start: int, byte_end: int):
        return _SegmentTraceContext(self, segment_id, byte_start, byte_end)

    def record_trace(self, trace: ExecutionTrace):
        self.traces.append(trace)
        if self.verbose:
            print(f"[TRACE] {trace.segment_id} [{trace.byte_start}:{trace.byte_end}] {trace.duration:.4f}s")

    def get_traces(self) -> List[ExecutionTrace]:
        return self.traces.copy()

    def reset(self):
        self.traces.clear()


class _SegmentTraceContext:
    def __init__(self, tracer: ExecutionTracer, segment_id: str, byte_start: int, byte_end: int):
        self.tracer = tracer
        self.segment_id = segment_id
        self.byte_start = byte_start
        self.byte_end = byte_end
        self.start_time = None
        self.exception = None

    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        duration = time.time() - self.start_time
        exception = None
        if exc_type:
            exception = f"{exc_type.__name__}: {exc_val}"

        trace = ExecutionTrace(
            segment_id=self.segment_id,
            byte_start=self.byte_start,
            byte_end=self.byte_end,
            duration=duration,
            exception=exception
        )
        self.tracer.record_trace(trace)
        return False