"""Data models for the RAG (Retrieval-Augmented Generation) bug-fix knowledge base."""

from __future__ import annotations

from pydantic import BaseModel, Field


class BugFixExample(BaseModel):
    """A single bug-fix pair from the dataset."""
    id: str
    dataset: str  # "BugsInPy" or "SWE-bench"
    repository: str
    example_kind: str = "full_example"  # full_example, function_snippet, diff_hunk, task_fix
    parent_example_id: str | None = None
    snippet_role: str | None = None
    snippet_index: int | None = None
    commit_hash: str | None = None
    issue_url: str | None = None

    # The bug description
    bug_description: str
    bug_summary: str  # Short one-line summary
    bug_category: str  # e.g., "off-by-one", "missing-null-check", "type-error"

    # The code change
    file_path: str
    function_name: str | None = None
    line_number: int | None = None

    # Before and after code snippets
    buggy_code: str
    fixed_code: str
    diff: str

    # Metadata for retrieval
    language: str = "Python"
    severity: str = "medium"  # low, medium, high, critical
    complexity: str = "simple"  # simple, moderate, complex

    # Embedding vector (stored separately in practice)
    embedding: list[float] | None = None

    # Additional context
    root_cause: str | None = None
    symptoms: list[str] = Field(default_factory=list)
    related_patterns: list[str] = Field(default_factory=list)
    task_mode: str | None = None
    validation_failures: list[str] = Field(default_factory=list)
    source_issue_types: list[str] = Field(default_factory=list)
    retrieval_tags: list[str] = Field(default_factory=list)


class RetrievalQuery(BaseModel):
    """Query for retrieving similar bug-fix examples."""
    query_code: str
    bug_description: str | None = None
    file_path: str | None = None
    function_name: str | None = None

    # Filters
    categories: list[str] | None = None
    severity: str | None = None
    error_classes: list[str] = Field(default_factory=list)
    preferred_categories: list[str] = Field(default_factory=list)
    preferred_datasets: list[str] = Field(default_factory=list)
    preferred_example_kinds: list[str] = Field(
        default_factory=lambda: ["function_snippet", "diff_hunk", "full_example"]
    )
    exclude_datasets: list[str] = Field(default_factory=list)
    required_terms: list[str] = Field(default_factory=list)
    max_per_dataset: int | None = 2
    max_results: int = 5


class RetrievedExample(BaseModel):
    """A retrieved bug-fix example with relevance score."""
    example: BugFixExample
    relevance_score: float
    match_reason: str  # Why this was matched


class RAGRetrievalResult(BaseModel):
    """Result from a RAG retrieval query."""
    query: str
    examples: list[RetrievedExample]
    total_found: int
    suggested_patterns: list[str]  # Patterns extracted from results
    context_summary: str  # A natural language summary of findings


class DatasetStatistics(BaseModel):
    """Statistics about the RAG knowledge base."""
    total_examples: int
    by_dataset: dict[str, int]
    by_category: dict[str, int]
    by_severity: dict[str, int]
    last_updated: str
