import chromadb
from sentence_transformers import SentenceTransformer
from flask import Flask, request, jsonify, send_from_directory

model = SentenceTransformer(
    "Snowflake/snowflake-arctic-embed-m-long",
    device="cpu",               # <-- the only line that changes
    trust_remote_code=True
)

# Test Test Test
client = chromadb.HttpClient(host="chromadb", port=8000)

collection = client.get_or_create_collection("news")


def query(question):

    query_embedding = model.encode(question, normalize_embeddings=True)

    results = collection.query(
        query_embeddings=query_embedding,
        n_results=5,
    )

    return results

app = Flask(__name__)

@app.route("/query", methods=["POST"])
def query_endpoint():
    data = request.get_json()
    question = data.get("question")
    if not question:
        return jsonify({"error": "Missing 'question' in request body"}), 400

    results = query(question)

    # Adapt response to MCP Model standards
    # Example MCP format:
    # {
    #   "results": [
    #       {
    #           "document": ...,
    #           "score": ...,
    #           "metadata": {...}
    #       },
    #       ...
    #   ]
    # }
    mcp_results = []
    for doc, score, meta in zip(
            results.get("documents", [[]])[0],
            results.get("distances", [[]])[0],
            results.get("metadatas", [[]])[0]):
        mcp_results.append({
            "document": doc,
            "score": float(score),
            "metadata": meta
        })

    return jsonify({"results": mcp_results})

@app.route("/health", methods=["GET"])
def health():
    return jsonify({"status": "ok"})

@app.route("/info", methods=["GET"])
def info():
    return jsonify({
        "provider": "StockDoc",
        "model": "Snowflake/snowflake-arctic-embed-m-long",
        "embedding_dim": 1024,  # or your actual dimension
        "collection": "news"
    })

@app.route("/tools", methods=["GET"])
def tools():
    return jsonify({
        "endpoints": [
            {
                "path": "/query",
                "method": "POST",
                "description": "Query the vector database with a question.",
                "request_format": {"question": "string"},
                "response_format": {
                    "results": [
                        {"document": "string", "score": "float", "metadata": "object"}
                    ]
                }
            },
            {
                "path": "/health",
                "method": "GET",
                "description": "Health check endpoint."
            },
            {
                "path": "/info",
                "method": "GET",
                "description": "Returns model and service metadata."
            }
        ]
    })

@app.route("/openapi.json", methods=["GET"])
def openapi():
    return send_from_directory('.', 'openapi.json')

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5005, threaded=True, debug=True)