"""Core data types for the neuromind.alzheimers sub-package."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Dict, List


class RAGMode(str, Enum):
    """RAG operating modes, ordered by increasing evidence breadth.

    Each mode maps to a fixed *top_k* retrieval count and a maximum
    context-token budget, mirroring the six conditions in the
    ``RAG_Comparison_Demo`` notebook.
    """

    NO_RAG = "no_rag"                  # LLM uses training knowledge only
    RETRIEVAL_ONLY = "retrieval_only"  # Raw excerpts shown, no generation
    BASIC = "basic"                    # 5 papers  (~20 K tokens)
    STANDARD = "standard"              # 10 papers (~40 K tokens)
    DETAILED = "detailed"              # 15 papers (~70 K tokens)
    COMPREHENSIVE = "comprehensive"    # 25 papers (~100 K tokens)


# Top-k documents retrieved per mode
RAG_MODE_TOP_K: Dict[RAGMode, int] = {
    RAGMode.NO_RAG: 0,
    RAGMode.RETRIEVAL_ONLY: 5,
    RAGMode.BASIC: 5,
    RAGMode.STANDARD: 10,
    RAGMode.DETAILED: 15,
    RAGMode.COMPREHENSIVE: 25,
}

# Approximate context-token budget per mode
RAG_MODE_MAX_TOKENS: Dict[RAGMode, int] = {
    RAGMode.NO_RAG: 0,
    RAGMode.RETRIEVAL_ONLY: 0,
    RAGMode.BASIC: 20_000,
    RAGMode.STANDARD: 40_000,
    RAGMode.DETAILED: 70_000,
    RAGMode.COMPREHENSIVE: 100_000,
}


@dataclass
class RetrievalResult:
    """Output of a single vector-store retrieval call."""

    documents: List[str]
    metadatas: List[Dict]
    similarities: List[float]
    retrieve_time: float  # seconds


@dataclass
class RAGResult:
    """Output of a complete RAG query (retrieval + generation)."""

    text: str
    mode: RAGMode
    query: str
    sources: List[Dict]
    context_tokens: int
    response_words: int
    retrieve_time: float   # seconds
    generate_time: float   # seconds

    @property
    def total_time(self) -> float:
        """Total wall-clock time for retrieval + generation."""
        return self.retrieve_time + self.generate_time

    def summary(self) -> str:
        """One-line human-readable summary."""
        return (
            f"[{self.mode.value}] {self.response_words} words | "
            f"{len(self.sources)} sources | "
            f"{self.context_tokens:,} ctx tokens | "
            f"{self.total_time:.1f}s"
        )


@dataclass
class IndexStats:
    """Statistics reported after an indexing run."""

    total_files: int = 0
    indexed: int = 0
    skipped: int = 0
    failed: int = 0
    total_chunks: int = 0

    @property
    def success_rate(self) -> float:
        if self.total_files == 0:
            return 0.0
        return self.indexed / self.total_files * 100

    def __str__(self) -> str:
        return (
            f"IndexStats(indexed={self.indexed}, skipped={self.skipped}, "
            f"failed={self.failed}, chunks={self.total_chunks:,})"
        )
