app.py

import json
import os
from decimal import Decimal

import psycopg2
import requests
from dotenv import load_dotenv
from flask import Flask, jsonify, request
from flask_cors import CORS

# Load environment variables from .env file
load_dotenv()

# Flask app initialization
app = Flask(__name__)
CORS(
    app,
    origins=["http://localhost:3000"],
    methods=["GET", "POST", "OPTIONS"],
    allow_headers=["Content-Type"],
    supports_credentials=True,
)
# Gemini 1.5 Flash API Configuration
API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent"
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
headers = {"Content-Type": "application/json"}

# Database Configuration
DB_CONFIG = {
    "dbname": os.getenv("DB_NAME"),
    "user": os.getenv("DB_USER"),
    "password": os.getenv("DB_PASSWORD"),
    "host": os.getenv("DB_HOST"),
    "port": os.getenv("DB_PORT"),
}

# Global variables to hold database connection
conn = None
cur = None


def get_connection():
    """Get the database connection and cursor."""
    global conn, cur
    if conn is None or cur is None:
        conn = psycopg2.connect(**DB_CONFIG)
        cur = conn.cursor()
    return conn, cur


def close_connection():
    """Close the database connection and cursor."""
    global conn, cur
    if cur:
        cur.close()
    if conn:
        conn.close()


def get_schema():
    """Fetch database schema."""
    try:
        conn, cur = get_connection()

        # Get all tables
        cur.execute(
            """
            SELECT table_name 
            FROM information_schema.tables 
            WHERE table_schema = 'public'
        """
        )
        tables = cur.fetchall()

        schema_text = "CREATE TABLE statements:\n\n"

        for table in tables:
            table_name = table[0]
            # Get columns and their details
            cur.execute(
                f"""
                SELECT 
                    c.column_name, 
                    c.data_type,
                    c.is_nullable,
                    c.column_default,
                    CASE 
                        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRIMARY KEY'
                        ELSE ''
                    END as key_type
                FROM information_schema.columns c
                LEFT JOIN information_schema.table_constraints tc
                    ON c.table_name = tc.table_name
                    AND tc.constraint_type = 'PRIMARY KEY'
                WHERE c.table_name = '{table_name}'
                ORDER BY c.ordinal_position;
            """
            )
            columns = cur.fetchall()

            # Create CREATE TABLE statement
            create_table = f"CREATE TABLE {table_name} (\n"
            for col in columns:
                name, data_type, nullable, default, key_type = col
                create_table += f"    {name} {data_type}"
                if key_type:
                    create_table += f" {key_type}"
                if nullable == "NO":
                    create_table += " NOT NULL"
                if default:
                    create_table += f" DEFAULT {default}"
                create_table += ",\n"
            create_table = create_table.rstrip(",\n") + "\n);\n\n"
            schema_text += create_table

        return schema_text

    except Exception as e:
        print(f"Error fetching schema: {e}")
        return None


def get_sql_query(user_query, schema):
    """Get SQL query from Gemini 1.5 Flash model."""
    try:
        prompt = f"""### PostgreSQL database schema:
{schema}

### Instructions:
- Generate a single, correct SQL query that answers the user's question.
- Include appropriate table aliases, JOIN conditions, and WHERE clauses.
- Provide the SQL query only, with no additional explanations or questions.

### User Query: {user_query}

### SQL Query:"""

        payload = {"contents": [{"parts": [{"text": prompt}]}]}

        response = requests.post(
            f"{API_URL}?key={GEMINI_API_KEY}", headers=headers, json=payload
        )

        if response.status_code != 200:
            print(f"Error from API: {response.text}")
            return None

        response_data = response.json()
        if "candidates" not in response_data or not response_data["candidates"]:
            print("No candidates in response")
            return None

        sql_query = response_data["candidates"][0]["content"]["parts"][0][
            "text"
        ].strip()
        if sql_query.startswith("```sql") and sql_query.endswith("```"):
            sql_query = sql_query[7:-3].strip()

        if not sql_query.endswith(";"):
            sql_query += ";"

        return sql_query

    except Exception as e:
        print(f"Error getting SQL query: {e}")
        return None


def execute_query(sql_query):
    """Execute SQL query and return results."""
    try:
        conn, cur = get_connection()
        cur.execute(sql_query)

        if sql_query.strip().upper().startswith("SELECT"):
            columns = [desc[0] for desc in cur.description]
            results = []
            for row in cur.fetchall():
                row_dict = {
                    columns[i]: float(row[i]) if isinstance(row[i], Decimal) else row[i]
                    for i in range(len(row))
                }
                results.append(row_dict)
        else:
            results = {"message": "Query executed successfully"}

        return results

    except Exception as e:
        print(f"Error executing query: {e}")
        return None


@app.route("/query", methods=["POST"])
def query():
    try:
        data = request.json
        user_query = data.get("query")
        if not user_query:
            return jsonify({"error": "Query is required"}), 400

        schema = get_schema()
        if not schema:
            return jsonify({"error": "Failed to fetch schema"}), 500

        sql_query = get_sql_query(user_query, schema)
        if not sql_query:
            return jsonify({"error": "Failed to generate SQL query"}), 500

        results = execute_query(sql_query)
        if results is None:
            return jsonify({"error": "Failed to execute query"}), 500

        return jsonify({"sql_query": sql_query, "results": results})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


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


if __name__ == "__main__":
    app.run(debug=True)