main.py

import os
import duckdb
from fastapi import FastAPI, HTTPException, Body
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel, Field
from pathlib import Path
import logging
import time # Import time for potential startup delays
import asyncio

# --- Configuration ---
DB_DIR = Path("data")
DB_FILENAME = "mydatabase.db"
DB_FILE = DB_DIR / DB_FILENAME
UI_EXPECTED_PORT = 8080 # Default port DuckDB UI often tries first

# Ensure the data directory exists
DB_DIR.mkdir(parents=True, exist_ok=True)

# --- Logging Setup ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- FastAPI App ---
app = FastAPI(
    title="DuckDB API & UI Host",
    description="Interact with DuckDB via API (/query, /download) and access the official DuckDB Web UI.",
    version="1.0.0"
)

# --- Pydantic Models ---
class QueryRequest(BaseModel):
    sql: str = Field(..., description="The SQL query to execute against DuckDB.")

class QueryResponse(BaseModel):
    columns: list[str] | None = None
    rows: list[dict] | None = None
    message: str | None = None
    error: str | None = None

# --- Helper Function ---
def execute_duckdb_query(sql_query: str, db_path: str = str(DB_FILE)):
    """Connects to DuckDB, executes a query, and returns results or error."""
    con = None
    try:
        logger.info(f"Connecting to database: {db_path}")
        con = duckdb.connect(database=db_path, read_only=False)
        logger.info(f"Executing SQL: {sql_query[:200]}{'...' if len(sql_query) > 200 else ''}")

        con.begin()
        result_relation = con.execute(sql_query)
        response_data = {"columns": None, "rows": None, "message": None, "error": None}

        if result_relation.description:
            columns = [desc[0] for desc in result_relation.description]
            rows_raw = result_relation.fetchall()
            rows_dict = [dict(zip(columns, row)) for row in rows_raw]
            response_data["columns"] = columns
            response_data["rows"] = rows_dict
            response_data["message"] = f"Query executed successfully. Fetched {len(rows_dict)} row(s)."
            logger.info(f"Query successful, returned {len(rows_dict)} rows.")
        else:
            response_data["message"] = "Query executed successfully (no data returned)."
            logger.info("Query successful (no data returned).")

        con.commit()
        return response_data

    except duckdb.Error as e:
        logger.error(f"DuckDB Error: {e}")
        if con: con.rollback()
        return {"columns": None, "rows": None, "message": None, "error": str(e)}
    except Exception as e:
        logger.error(f"General Error: {e}")
        if con: con.rollback()
        return {"columns": None, "rows": None, "message": None, "error": f"An unexpected error occurred: {e}"}
    finally:
        if con:
            con.close()
            logger.info("Database connection closed.")


# --- FastAPI Startup Event ---
@app.on_event("startup")
async def startup_event():
    logger.info("Application startup: Initializing DuckDB UI...")
    con = None
    try:
        # Connect to the main DB file to execute initialization commands
        # Use a temporary in-memory DB for UI start if main DB doesn't exist yet?
        # No, start_ui seems to need the target DB. Ensure DB file path exists.
        if not DB_FILE.parent.exists():
             DB_FILE.parent.mkdir(parents=True, exist_ok=True)

        # It's crucial the UI extension can write its state.
        # By default it uses ~/.duckdb/ which will be /root/.duckdb in the container.
        # Ensure this is writable or mount a volume there.
        logger.info(f"Attempting to connect to {DB_FILE} for UI setup.")
        con = duckdb.connect(database=str(DB_FILE), read_only=False)

        logger.info("Installing and loading 'ui' extension...")
        con.execute("INSTALL ui;")
        con.execute("LOAD ui;")

        logger.info("Calling start_ui()... This will start a separate web server.")
        # CALL start_ui() starts the server in the background (usually)
        # It might print the URL/port it's using to stderr/stdout of the main process
        con.execute("CALL start_ui();")

        # Give the UI server a moment to start up. This is a guess.
        # A more robust solution might involve checking if the port is listening.
        await asyncio.sleep(2)

        logger.info(f"DuckDB UI server startup initiated. It usually listens on port {UI_EXPECTED_PORT}.")
        logger.info("Check container logs for the exact URL if it differs.")
        logger.info("API server (FastAPI/Uvicorn) is running on port 8000.")

    except duckdb.Error as e:
        logger.error(f"CRITICAL: Failed to install/load/start DuckDB UI extension: {e}")
        logger.error("The DuckDB UI will likely not be available.")
    except Exception as e:
        logger.error(f"CRITICAL: An unexpected error occurred during UI startup: {e}")
        logger.error("The DuckDB UI will likely not be available.")
    finally:
        if con:
            con.close()
            logger.info("UI setup connection closed.")

# --- API Endpoints ---
@app.get("/", summary="Root Endpoint / Info", tags=["General"])
async def read_root():
    """Provides links to the API docs and the DuckDB UI."""
    # Assumes UI is running on localhost from the container's perspective
    # User needs to map the port correctly
    return JSONResponse({
        "message": "DuckDB API and UI Host",
        "api_details": {
            "docs": "/docs",
            "query_endpoint": "/query (POST)",
            "download_endpoint": "/download (GET)"
        },
        "duckdb_ui": {
            "message": f"Access the official DuckDB Web UI. It should be running on port {UI_EXPECTED_PORT} inside the container.",
            "typical_access_url": f"http://localhost:{UI_EXPECTED_PORT}",
            "notes": f"Ensure you have mapped port {UI_EXPECTED_PORT} from the container when running `docker run` (e.g., -p {UI_EXPECTED_PORT}:{UI_EXPECTED_PORT})."
        },
        "database_file_container_path": str(DB_FILE)
    })

@app.post("/query", response_model=QueryResponse, summary="Execute SQL Query", tags=["Database API"])
async def execute_query_endpoint(query_request: QueryRequest):
    """
    Executes a given SQL query against the DuckDB database via the API.
    Handles SELECT, INSERT, UPDATE, DELETE, CREATE TABLE, etc.
    """
    result = execute_duckdb_query(query_request.sql)
    if result["error"]:
        raise HTTPException(status_code=400, detail=result["error"])
    return JSONResponse(content=result)


@app.get("/download", summary="Download Database File", tags=["Database API"])
async def download_database_file():
    """
    Allows downloading the current DuckDB database file via the API.
    """
    if not DB_FILE.is_file():
        logger.error(f"Download request failed: Database file not found at {DB_FILE}")
        raise HTTPException(status_code=404, detail="Database file not found.")

    logger.info(f"Serving database file for download: {DB_FILE}")
    return FileResponse(
        path=str(DB_FILE),
        filename=DB_FILENAME,
        media_type='application/octet-stream'
    )

# Need asyncio for sleep in startup
# import asyncio

# --- Run with Uvicorn (for local testing - doesn't handle UI startup well here) ---
# if __name__ == "__main__":
#     # Note: Running directly with python main.py won't trigger the startup
#     # event correctly in the same way uvicorn command does.
#     # Use `uvicorn main:app --reload --port 8000` for local dev testing.
#     print("Run using: uvicorn main:app --host 0.0.0.0 --port 8000")