"""
Query schemas for the sim-search API.

This module defines the Pydantic models for query-related operations.
"""

from typing import Dict, List, Optional, Any

from pydantic import BaseModel


class QueryBase(BaseModel):
    """Base query schema."""
    
    query: str


class QueryProcess(QueryBase):
    """Query processing schema."""
    
    pass


class QueryClassify(QueryBase):
    """Query classification schema."""
    
    pass


class SubQuestion(BaseModel):
    """Sub-question schema."""
    
    sub_question: str
    aspect: str
    priority: float


class StructuredQuery(BaseModel):
    """Structured query schema."""
    
    original_query: str
    enhanced_query: Optional[str] = None
    type: Optional[str] = None
    intent: Optional[str] = None
    domain: Optional[str] = None
    confidence: Optional[float] = None
    reasoning: Optional[str] = None
    entities: Optional[List[str]] = None
    sub_questions: Optional[List[SubQuestion]] = None
    search_queries: Optional[Dict[str, str]] = None
    is_academic: Optional[bool] = None
    is_code: Optional[bool] = None
    is_current_events: Optional[bool] = None


class ProcessedQuery(BaseModel):
    """Processed query schema."""
    
    original_query: str
    structured_query: StructuredQuery
    
    class Config:
        """Pydantic config."""
        
        json_schema_extra = {
            "example": {
                "original_query": "What are the latest advancements in quantum computing?",
                "structured_query": {
                    "original_query": "What are the latest advancements in quantum computing?",
                    "enhanced_query": "What are the recent breakthroughs and developments in quantum computing technology, algorithms, and applications in the past 2 years?",
                    "type": "exploratory",
                    "intent": "research",
                    "domain": "academic",
                    "confidence": 0.95,
                    "reasoning": "This query is asking about recent developments in a scientific field, which is typical of academic research.",
                    "entities": ["quantum computing", "advancements"],
                    "sub_questions": [
                        {
                            "sub_question": "What are the latest hardware advancements in quantum computing?",
                            "aspect": "hardware",
                            "priority": 0.9
                        },
                        {
                            "sub_question": "What are the recent algorithmic breakthroughs in quantum computing?",
                            "aspect": "algorithms",
                            "priority": 0.8
                        }
                    ],
                    "search_queries": {
                        "google": "latest advancements in quantum computing 2024",
                        "scholar": "recent quantum computing breakthroughs",
                        "arxiv": "quantum computing hardware algorithms"
                    },
                    "is_academic": True,
                    "is_code": False,
                    "is_current_events": False
                }
            }
        }