import time
from typing import Dict, Optional
from dataclasses import dataclass


@dataclass
class ProfileSnapshot:
    elapsed: float
    call_count: int
    memory_approx: int = 0


class SegmentProfiler:
    def __init__(self):
        self._profiles: Dict[str, ProfileSnapshot] = {}
        self._active: Dict[str, float] = {}
        self._call_counts: Dict[str, int] = {}

    def start(self, segment_id: str):
        if segment_id in self._active:
            raise RuntimeError(f"Segment {segment_id} already being profiled")
        self._active[segment_id] = time.time()
        self._call_counts[segment_id] = self._call_counts.get(segment_id, 0) + 1

    def stop(self, segment_id: str):
        if segment_id not in self._active:
            raise RuntimeError(f"Segment {segment_id} not being profiled")

        elapsed = time.time() - self._active.pop(segment_id)
        count = self._call_counts[segment_id]

        if segment_id in self._profiles:
            prev = self._profiles[segment_id]
            new_elapsed = prev.elapsed + elapsed
            new_count = prev.call_count + count
            self._profiles[segment_id] = ProfileSnapshot(
                elapsed=new_elapsed,
                call_count=new_count
            )
        else:
            self._profiles[segment_id] = ProfileSnapshot(
                elapsed=elapsed,
                call_count=count
            )

    def get_profile(self, segment_id: str) -> Optional[ProfileSnapshot]:
        return self._profiles.get(segment_id)

    def get_all_profiles(self) -> Dict[str, ProfileSnapshot]:
        return self._profiles.copy()

    def reset(self):
        self._profiles.clear()
        self._active.clear()
        self._call_counts.clear()