from fastapi import FastAPI, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Dict, Optional, List, Any
import uuid
from datetime import datetime
from contextlib import asynccontextmanager

from models.embedding import EmbeddingModel
from models.summarization import SummarizationModel
from models.nlp import NLPModel
from database.query import DatabaseService
from database.query_processor import QueryProcessor

# Initialize models
embedding_model = None
summarization_model = None
nlp_model = None
db_service = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load models when app starts
    global embedding_model, summarization_model, nlp_model, db_service
    embedding_model = EmbeddingModel()
    summarization_model = SummarizationModel()
    nlp_model = NLPModel()
    db_service = DatabaseService()
    yield
    # Clean up when app stops
    if db_service:
        await db_service.close()

app = FastAPI(
    title="Kairos News API",
    version="1.0",
    lifespan=lifespan
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# In-memory job storage
jobs_db: Dict[str, Dict] = {}

class PostRequest(BaseModel):
    query: str
    topic: Optional[str] = None
    start_date: Optional[str] = None  # Format: "YYYY-MM-DD"
    end_date: Optional[str] = None    # Format: "YYYY-MM-DD"

class JobStatus(BaseModel):
    id: str
    status: str  # "processing", "completed", "failed"
    created_at: datetime
    completed_at: Optional[datetime] = None
    request: PostRequest
    result: Optional[Dict[str, Any]] = None  # Flexible result structure

@app.post("/index", response_model=JobStatus)
async def create_job(request: PostRequest, background_tasks: BackgroundTasks):
    job_id = str(uuid.uuid4())
    
    jobs_db[job_id] = {
        "id": job_id,  # Ensure `id` is included
        "status": "processing",
        "created_at": datetime.now(),
        "completed_at": None,
        "request": request.dict(),
        "result": None
    }

    background_tasks.add_task(
        process_job,
        job_id,
        request,
        embedding_model,
        summarization_model,
        nlp_model,
        db_service
    )
    
    return jobs_db[job_id]  # Return the full job object

@app.get("/loading", response_model=JobStatus)
async def get_job_status(id: str):
    if id not in jobs_db:
        raise HTTPException(status_code=404, detail="Job not found")
    
    return jobs_db[id]

async def process_job(
    job_id: str,
    request: PostRequest,
    embedding_model: EmbeddingModel,
    summarization_model: SummarizationModel,
    nlp_model: NLPModel,
    db_service: DatabaseService
):
    try:
        processor = QueryProcessor(
            embedding_model=embedding_model,
            summarization_model=summarization_model,
            nlp_model=nlp_model,
            db_service=db_service
        )
        
        result = await processor.process(
            query=request.query,
            topic=request.topic,
            start_date=request.start_date,
            end_date=request.end_date
        )
        
        jobs_db[job_id].update({
            "status": "completed",
            "completed_at": datetime.now(),
            "result": result if result else {"message": "No results found"}
        })
    except Exception as e:
        jobs_db[job_id].update({
            "status": "failed",
            "completed_at": datetime.now(),
            "result": {"error": str(e)}
        })