from fastapi import FastAPI, Depends, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.orm import Session
from sqlalchemy import text
from database import get_db, init_db, Email, EMBEDDING_DIMENSIONS
from pydantic import BaseModel
from typing import List, Optional
from google import genai
from google.genai import types
import os
import time

app = FastAPI(title="Unified Email Semantic Search API")

# Setup CORS for the SPA
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize Gemini client — reads GEMINI_API_KEY from environment
gemini_client = None

class SearchQuery(BaseModel):
    query: str
    limit: int = 10

class SearchResult(BaseModel):
    message_id: str
    subject: str
    sender: str
    date: str
    snippet: str
    distance: float

@app.on_event("startup")
def on_startup():
    global gemini_client
    print("Initializing Database...")
    time.sleep(2) # Give postgres a moment to be fully ready
    try:
        init_db()
    except Exception as e:
        print(f"Error initializing DB: {e}")

    api_key = os.environ.get("GEMINI_API_KEY")
    if api_key:
        gemini_client = genai.Client(api_key=api_key)
        print("Gemini client initialized.")
    else:
        print("WARNING: GEMINI_API_KEY not set. Embedding features disabled.")

@app.post("/search", response_model=List[SearchResult])
def search_emails(request: SearchQuery, db: Session = Depends(get_db)):
    if not gemini_client:
        raise HTTPException(status_code=500, detail="Gemini API Key is not configured.")

    try:
        response = gemini_client.models.embed_content(
            model="gemini-embedding-001",
            contents=request.query,
            config=types.EmbedContentConfig(
                task_type="RETRIEVAL_QUERY",
                output_dimensionality=EMBEDDING_DIMENSIONS,
            ),
        )
        query_embedding = response.embeddings[0].values
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Embedding API error: {e}")

    # Use pgvector's cosine distance operator via SQLAlchemy ORM
    results = db.query(
        Email, 
        Email.embedding.cosine_distance(query_embedding).label('distance')
    ).order_by(
        Email.embedding.cosine_distance(query_embedding)
    ).limit(request.limit).all()

    response_data = []
    for email, distance in results:
        # Create a snippet from the content
        snippet = email.content[:200] + "..." if email.content and len(email.content) > 200 else (email.content or "")
        response_data.append(SearchResult(
            message_id=email.message_id or "",
            subject=email.subject or "",
            sender=email.sender or "",
            date=email.date.isoformat() if email.date else "",
            snippet=snippet,
            distance=distance
        ))

    return response_data