from flask import Blueprint, request, jsonify, Response
from train import train_vectorstore
from query import query_vectorstore
import uuid

import os
from werkzeug.utils import secure_filename

bp = Blueprint("api", __name__)

ALLOWED_EXTENSIONS = {"pdf", "md"}

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@bp.route("/upload", methods=["POST"])
def upload_file():
    if "file" not in request.files:
        return jsonify({"status": "error", "message": "No file part in the request."}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"status": "error", "message": "No file selected."}), 400

    if file and allowed_file(file.filename):
        upload_folder = "./data"
        os.makedirs(upload_folder, exist_ok=True)

        # Extract the extension and generate UUID-based filename
        ext = file.filename.rsplit('.', 1)[1].lower()
        unique_filename = f"{uuid.uuid4()}.{ext}"
        secure_unique_filename = secure_filename(unique_filename)

        file_path = os.path.join(upload_folder, secure_unique_filename)
        file.save(file_path)

        return jsonify({
            "status": "success",
            "message": "File uploaded successfully.",
            "file": secure_unique_filename
        }), 200

    return jsonify({"status": "error", "message": "File type not allowed."}), 400

@bp.route("/", methods=["GET"])
def documentation():
    markdown_content = """
# API Documentation

Welcome to the API Documentation for our vector store system. Below are the available endpoints.

---

## POST /api/train

**Description:**  
Trains the vector store by loading documents from various sources.

**Expected JSON Payload:**

- **type**: `"file"`, `"website"`, `"text"`, or `"youtube"`
  - For **file**: provide `"file_path"` (string or list of strings)
  - For **website**: provide `"urls"` (list of URL strings)
  - For **text**: provide `"text_contents"` (list of strings)
  - For **youtube**: provide `"youtube_links"` (list of YouTube URL strings)
- **query** (optional): a test query to execute against the trained index

**Response:**  
A JSON object with a `status` and `message`, and if a test query is provided, a `test_response`.

---

## GET /api/query

**Description:**  
Queries the pre-trained vector store.

**Query Parameters:**

- `query` (required): the query text to search the index.

**Response:**  
A JSON object with a `status` and the `response` from the query.

---

Enjoy using the API!
    """
    return Response(markdown_content, status=200, mimetype="text/markdown")

@bp.route("/train", methods=["POST"])
def train_endpoint():
    try:
        result, status_code = train_vectorstore()
        return result, status_code
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


@bp.route("/query", methods=["POST"])
def query_endpoint():

    try:
        result = query_vectorstore()
        return jsonify(result), 200
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


