"""Data models for retrieval learning."""
from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any


def utcnow() -> str:
    """Return current UTC timestamp as ISO format string."""
    return datetime.now(UTC).isoformat()


@dataclass
class RetrievalSettings:
    """Learned retrieval settings for a query segment."""

    limit: int
    filter: dict[str, Any] | None = None
    suggested_chunk_size: int | None = None
    suggested_overlap: int | None = None


@dataclass
class QueryRewriteRecord:
    """Successful or failed query rewrite accumulated from prior runs."""

    original_query: str
    rewritten_query: str
    workspace_id: str | None = None
    task_type: str | None = None
    user_segment: str | None = None
    success_count: int = 0
    failure_count: int = 0
    last_used_at: str = field(default_factory=utcnow)


@dataclass
class RetrievalEvent:
    """Recorded retrieval call linked to a trace."""

    id: str
    trace_id: str | None
    session_id: str | None
    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    prompt_name: str | None
    prompt_version: str | None
    original_query: str
    effective_query: str
    limit: int
    filter: dict[str, Any] | None
    chunks: list[dict[str, Any]]
    created_at: str = field(default_factory=utcnow)
    quality_score: float = 0.0
    bad_context_detected: bool = False
    synced_at: str | None = None
    applied_feedback_ids: list[str] = field(default_factory=list)
    applied_outcome_ids: list[str] = field(default_factory=list)


@dataclass
class SourceUsefulnessScore:
    """Aggregated usefulness signal for a source or connector."""

    workspace_id: str | None
    source_key: str
    connector: str | None = None
    helpful_count: int = 0
    harmful_count: int = 0
    accepted_count: int = 0
    rejected_count: int = 0
    usefulness_score: float = 0.0
    last_updated: str = field(default_factory=utcnow)


@dataclass
class ReindexTrigger:
    """Signal that a source should be re-indexed or re-chunked."""

    id: str
    workspace_id: str | None
    source_key: str
    reason: str
    severity: str
    occurrences: int = 1
    status: str = "open"
    suggested_chunk_size: int | None = None
    suggested_overlap: int | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=utcnow)
    updated_at: str = field(default_factory=utcnow)


@dataclass
class ReindexRequest:
    """Manually requested reindex operation."""

    id: str
    workspace_id: str | None
    source_key: str
    reason: str
    requested_by: str | None = None
    status: str = "queued"
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: str = field(default_factory=utcnow)
    updated_at: str = field(default_factory=utcnow)


@dataclass
class SourceDocumentSnapshot:
    """Stored document snapshot used by the reindex worker."""

    workspace_id: str | None
    source_key: str
    document_id: str
    content: str
    metadata: dict[str, Any]
    source_ref: str | None = None
    updated_at: str = field(default_factory=utcnow)


@dataclass
class RetrievalProfile:
    """Aggregated retrieval profile for a segment."""

    workspace_id: str | None
    task_type: str | None
    user_segment: str | None
    limit: int
    filter: dict[str, Any]
    success_count: int = 0
    failure_count: int = 0
    score: float = 0.0
    last_updated: str = field(default_factory=utcnow)
