"""Audit record storage with filtering and retrieval.

In production, this would be backed by PostgreSQL.
This implementation provides the storage interface for testing.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from uuid import UUID

from services.intelligence_pipeline_v3.audit.models import (
    AuditRecord,
    CorrectionEvent,
    ReviewFilter,
)


@dataclass
class AuditStore:
    """In-memory audit record store with filtering.

    Provides storage, retrieval, and filtering of audit records and
    their immutable correction events.
    """

    _records: dict[UUID, AuditRecord] = field(default_factory=dict)
    _corrections: list[CorrectionEvent] = field(default_factory=list)

    def store(self, record: AuditRecord) -> None:
        """Store an audit record."""
        self._records[record.record_id] = record

    def get(self, record_id: UUID) -> AuditRecord | None:
        """Retrieve a record by ID."""
        return self._records.get(record_id)

    def get_by_document(self, document_id: str) -> list[AuditRecord]:
        """Get all records for a document."""
        return [
            r for r in self._records.values() if r.document_id == document_id
        ]

    def get_by_run(self, run_id: UUID) -> AuditRecord | None:
        """Get the record for a pipeline run."""
        for r in self._records.values():
            if r.run_id == run_id:
                return r
        return None

    def add_correction(
        self, record_id: UUID, correction: CorrectionEvent
    ) -> bool:
        """Add a correction to a record. Returns False if record not found."""
        record = self._records.get(record_id)
        if record is None:
            return False
        record.add_correction(correction)
        self._corrections.append(correction)
        return True

    def filter(self, criteria: ReviewFilter) -> list[AuditRecord]:
        """Filter records by criteria."""
        return [
            r for r in self._records.values() if criteria.matches(r)
        ]

    def get_corrections(self, record_id: UUID) -> list[CorrectionEvent]:
        """Get all corrections for a record."""
        record = self._records.get(record_id)
        if record is None:
            return []
        return list(record.corrections)

    def count(self) -> int:
        """Total stored records."""
        return len(self._records)

    def correction_count(self) -> int:
        """Total correction events across all records."""
        return len(self._corrections)