from mcp.server.fastmcp import FastMCP
import sqlite3
import json

# Initialize MCP server
mcp = FastMCP("EnterpriseData")

# Create in-memory SQLite database
conn = sqlite3.connect(":memory:", check_same_thread=False)
cursor = conn.cursor()

# Create Customers table
cursor.execute("""
CREATE TABLE Customers (
    CustomerID INTEGER PRIMARY KEY AUTOINCREMENT,
    Name TEXT,
    Region TEXT,
    LastOrderDate TEXT
)
""")

# Seed sample data
cursor.executemany(
    "INSERT INTO Customers (Name, Region, LastOrderDate) VALUES (?, ?, ?)",
    [
        ("Acme Corp", "Northeast", "2024-12-01"),
        ("Beta Inc", "West", "2025-06-01"),
        ("Gamma Co", "Northeast", "2023-09-15"),
        ("Delta LLC", "South", "2025-03-20"),
        ("Epsilon Ltd", "Northeast", "2025-07-10")
    ]
)
conn.commit()

# Expose SQL tool to the AI
@mcp.tool()
def query_database(sql: str) -> str:
    """Executes raw SQL and returns JSON results."""
    try:
        cur = conn.cursor()
        cur.execute(sql)
        rows = cur.fetchall()
        columns = [desc[0] for desc in cur.description or []]
        results = [dict(zip(columns, row)) for row in rows]
        return json.dumps(results)
    except Exception as e:
        return json.dumps({"error": str(e)})

# Start MCP server
if __name__ == "__main__":
    mcp.run(transport="stdio")