#!/usr/bin/env python3
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import logging
import time
import json
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

STATUS_FILE = Path('/tmp/setup_status.json')
READY_FLAG = Path('/tmp/faiss_ready')

def get_setup_status():
    if not STATUS_FILE.exists():
        return {'status': 'initializing', 'message': 'Setup não iniciado', 'progress': 0}
    try:
        with open(STATUS_FILE) as f:
            return json.load(f)
    except:
        return {'status': 'unknown', 'message': 'Erro ao ler status', 'progress': 0}

def is_ready():
    return READY_FLAG.exists()

query_engine = None

def get_query_engine():
    global query_engine
    if query_engine is None:
        if not is_ready():
            raise HTTPException(status_code=503, detail="RAG em construção. Tente em alguns minutos.")
        logger.info("Carregando QueryEngine...")
        from query_engine import QueryEngine
        query_engine = QueryEngine()
        logger.info("✅ QueryEngine carregado!")
    return query_engine

app = FastAPI(title="Para.AI RAG Cluster (LangChain)", version="1.0.0")

class EmbeddingSearchRequest(BaseModel):
    query: str
    top_k: int = 10
    return_embeddings: bool = False

class KeywordSearchRequest(BaseModel):
    keywords: List[str]
    operator: str = "AND"
    top_k: int = 20

class IDSearchRequest(BaseModel):
    ids: List[str]
    return_embeddings: bool = False

@app.get("/")
async def root():
    setup_status = get_setup_status()
    ready = is_ready()

    response = {"status": "online", "rag_ready": ready, "setup": setup_status, "backend": "LangChain + FAISS (CPU)"}

    if ready and query_engine:
        response["cluster_id"] = query_engine.config.get('cluster_id')
        response["chunk_range"] = [query_engine.config.get('chunk_start'), query_engine.config.get('chunk_end')]

    return response

@app.get("/setup/status")
async def setup_status():
    return get_setup_status()

@app.get("/health")
async def health():
    return {"status": "ok", "timestamp": time.time()}

@app.post("/search/embedding")
async def search_embedding(request: EmbeddingSearchRequest):
    engine = get_query_engine()
    try:
        start = time.time()
        results = engine.search_by_embedding(request.query, request.top_k, request.return_embeddings)
        results['query_time_ms'] = round((time.time() - start) * 1000, 2)
        return results
    except Exception as e:
        logger.error(f"Erro: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/search/keywords")
async def search_keywords(request: KeywordSearchRequest):
    engine = get_query_engine()
    try:
        start = time.time()
        results = engine.search_by_keywords(request.keywords, request.operator, request.top_k)
        results['query_time_ms'] = round((time.time() - start) * 1000, 2)
        return results
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/search/by_id")
async def search_by_id(request: IDSearchRequest):
    engine = get_query_engine()
    try:
        start = time.time()
        results = engine.search_by_ids(request.ids, request.return_embeddings)
        results['query_time_ms'] = round((time.time() - start) * 1000, 2)
        return results
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/cluster/info")
async def cluster_info():
    engine = get_query_engine()
    try:
        info = engine.get_cluster_info()
        info['uptime_seconds'] = round(time.time() - app.state.start_time, 2)
        return info
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.on_event("startup")
async def startup_event():
    app.state.start_time = time.time()
    logger.info("="*80)
    logger.info("🚀 Para.AI RAG (LangChain + FAISS) ONLINE")
    logger.info("="*80)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)