"""
Voyage AI RAG API (Flask)
Generated by vai v{{vaiVersion}} on {{generatedAt}}

Endpoints:
- POST /api/search  - Semantic search with optional reranking
- POST /api/ingest  - Document ingestion
- GET  /api/health  - Health check
"""

import os
import time
from flask import Flask, request, jsonify
from dotenv import load_dotenv

load_dotenv()

from voyage_client import embed_query{{#if rerank}}, rerank{{/if}}
from mongo_client import connect, vector_search, insert_documents, close
from chunker import chunk_text

app = Flask(__name__)


@app.route("/")
def index():
    """API information."""
    return jsonify({
        "name": "{{projectName}}",
        "description": "Voyage AI RAG API",
        "model": "{{model}}",
        "database": "{{db}}",
        "collection": "{{collection}}",
        "endpoints": {
            "search": "POST /api/search",
            "ingest": "POST /api/ingest",
            "health": "GET /api/health",
        },
    })


@app.route("/api/search", methods=["POST"])
def search():
    """
    Semantic search endpoint.
    
    Request body:
    {
        "query": "What is vector search?",
        "limit": 5,
        "include_context": true
    }
    """
    start = time.time()
    data = request.get_json()
    
    query = data.get("query")
    if not query:
        return jsonify({"error": "Query is required"}), 400
    
    limit = data.get("limit", 5)
{{#if rerank}}
    candidates = data.get("candidates", 20)
{{/if}}
    include_context = data.get("include_context", False)
    filter_query = data.get("filter")
    
    # Step 1: Embed the query
    query_embedding = embed_query(query)
    
    # Step 2: Vector search
{{#if rerank}}
    search_results = vector_search(query_embedding, limit=candidates, filter_query=filter_query)
{{else}}
    search_results = vector_search(query_embedding, limit=limit, filter_query=filter_query)
{{/if}}
    
    if not search_results:
        return jsonify({
            "results": [],
            "meta": {"model": "{{model}}", "took": int((time.time() - start) * 1000)},
        })
    
{{#if rerank}}
    # Step 3: Rerank for better relevance
    documents = [r["document"]["text"] for r in search_results]
    rerank_result = rerank(query, documents, top_k=limit)
    
    results = [
        {
            "text": search_results[r["index"]]["document"]["text"],
            "score": r["relevance_score"],
            "metadata": search_results[r["index"]]["document"].get("metadata", {}),
        }
        for r in rerank_result["results"]
    ]
{{else}}
    results = [
        {
            "text": r["document"]["text"],
            "score": r["score"],
            "metadata": r["document"].get("metadata", {}),
        }
        for r in search_results
    ]
{{/if}}
    
    response = {
        "results": results,
        "meta": {
            "model": "{{model}}",
{{#if rerank}}
            "rerank_model": "{{rerankModel}}",
{{/if}}
            "took": int((time.time() - start) * 1000),
        },
    }
    
    if include_context:
        response["context"] = "\n\n---\n\n".join(r["text"] for r in results)
    
    return jsonify(response)


@app.route("/api/ingest", methods=["POST"])
def ingest():
    """
    Document ingestion endpoint.
    
    Request body:
    {
        "text": "Document content...",
        "metadata": {"source": "api"}
    }
    """
    start = time.time()
    data = request.get_json()
    
    text = data.get("text")
    if not text:
        return jsonify({"error": "Text is required"}), 400
    
    metadata = data.get("metadata", {})
    metadata["source"] = metadata.get("source", "api")
    
    # Chunk the text
    chunks = chunk_text(text)
    
    # Embed all chunks
    from voyage_client import embed
    result = embed(chunks, input_type="document")
    embeddings = result["embeddings"]
    
    # Prepare documents
    documents = [
        {
            "text": chunk,
            "embedding": embeddings[i],
            "metadata": {
                **metadata,
                "chunk_index": i,
                "total_chunks": len(chunks),
            },
        }
        for i, chunk in enumerate(chunks)
    ]
    
    # Insert into MongoDB
    insert_result = insert_documents(documents)
    
    return jsonify({
        "success": True,
        "chunks": len(chunks),
        "tokens": result["usage"]["total_tokens"],
        "took": int((time.time() - start) * 1000),
    })


@app.route("/api/health", methods=["GET"])
def health():
    """Health check endpoint."""
    try:
        connect()
        return jsonify({
            "status": "healthy",
            "model": "{{model}}",
            "database": "{{db}}",
            "collection": "{{collection}}",
        })
    except Exception as e:
        return jsonify({
            "status": "unhealthy",
            "error": str(e),
        }), 503


@app.teardown_appcontext
def shutdown(exception=None):
    """Clean up on shutdown."""
    pass  # Connection pooling handles this


if __name__ == "__main__":
    port = int(os.getenv("PORT", 3000))
    print(f"""
🚀 Voyage AI RAG Server (Flask)

   Model:      {{model}}
   Database:   {{db}}.{{collection}}
{{#if rerank}}
   Reranking:  {{rerankModel}}
{{/if}}

   Endpoints:
   POST /api/search  - Semantic search
   POST /api/ingest  - Document ingestion
   GET  /api/health  - Health check

   Running on http://localhost:{port}
""")
    app.run(host="0.0.0.0", port=port, debug=os.getenv("FLASK_DEBUG", "0") == "1")