from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from src.search_engine import PromptSearchEngine
from src.prompt_loader import PromptLoader

# Constants
SEED = 42
DATA_SIZE = 100

# Initialize the prompt loader and search engine
prompts = PromptLoader(seed=SEED).load_data(size=DATA_SIZE)
engine = PromptSearchEngine(prompts)

# Initialize FastAPI
app = FastAPI()


# Request and Response Models
class QueryRequest(BaseModel):
    query: str
    n: int = 5


class SimilarPrompt(BaseModel):
    score: float
    prompt: str


class QueryResponse(BaseModel):
    similar_prompts: List[SimilarPrompt]


# API endpoint
@app.post("/most_similar", response_model=QueryResponse)
async def get_most_similar(query_request: QueryRequest):
    try:
        similar_prompts = engine.most_similar(
            query=query_request.query, n=query_request.n
        )
        response = QueryResponse(
            similar_prompts=[
                SimilarPrompt(score=score, prompt=prompt)
                for score, prompt in similar_prompts
            ]
        )
        return response
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))