import time
import uuid

class TraceContext:
  def __init__(self, manager, name, input=None):
    self.manager = manager
    self.name = name
    self.input = input
    self.output = None
    self.trace_id = str(uuid.uuid4())
    self.start_time = None

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

  def __exit__(self, exc_type, exc, tb):
    end_time = time.time()

    self.manager.push_all(
      trace_id=self.trace_id,
      name=self.name,
      input=self.input,
      output=self.output,
      start_time=self.start_time,
      end_time=end_time,
      error=str(exc) if exc else None
    )