"""
Repository for trace data storage and retrieval.
Thread-safe in-memory cache with O(1) lookups.
"""
from typing import Dict, List, Optional

from traces_explorer.backend.models.trace import Trace
from traces_explorer.backend.utils.rwlock import RWLock


class TraceRepository:
    """
    Thread-safe in-memory repository for traces.
    Uses RWLock for concurrent read access and exclusive writes.
    """

    def __init__(self):
        self._traces: Dict[str, Trace] = {}
        self._rwlock = RWLock()

    def add(self, trace: Trace) -> None:
        """
        Add a trace to the repository.

        Args:
            trace: Trace to add
        """
        with self._rwlock.write():
            self._traces[trace.id] = trace

    def add_all(self, traces: List[Trace]) -> None:
        """
        Add multiple traces in a single write lock.
        More efficient than multiple add() calls.

        Args:
            traces: List of Trace objects
        """
        with self._rwlock.write():
            for trace in traces:
                self._traces[trace.id] = trace

    def get_by_id(self, trace_id: str) -> Optional[Trace]:
        """
        Get a trace by ID. Thread-safe read.

        Args:
            trace_id: Trace ID to look up

        Returns:
            Trace if found, None otherwise
        """
        with self._rwlock.read():
            return self._traces.get(trace_id)

    def list_all(self) -> List[Trace]:
        """
        List all traces. Thread-safe read.

        Returns:
            List of Trace objects
        """
        with self._rwlock.read():
            return list(self._traces.values())

    def clear(self) -> None:
        """
        Clear all traces. Thread-safe write.
        """
        with self._rwlock.write():
            self._traces.clear()

    def remove(self, trace_id: str) -> None:
        """
        Remove a trace by ID if present. Thread-safe write.

        Args:
            trace_id: Trace ID to remove
        """
        with self._rwlock.write():
            self._traces.pop(trace_id, None)

    def remove_by_source(self, source: str) -> None:
        """
        Remove all traces for a given source. Thread-safe write.

        Args:
            source: Trace source value (e.g. "dataset", "pasted")
        """
        with self._rwlock.write():
            ids_to_remove = [
                trace_id
                for trace_id, trace in self._traces.items()
                if getattr(trace, "source", None) == source
            ]
            for trace_id in ids_to_remove:
                self._traces.pop(trace_id, None)

    def count(self) -> int:
        """
        Get count of traces. Thread-safe read.

        Returns:
            Number of traces in repository
        """
        with self._rwlock.read():
            return len(self._traces)
