# database_server.py
import asyncio
import sqlite3
from mcp.server import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server
import mcp.types as types

# Create the MCP server
app = Server("database-server")

# Connect to SQLite (replace with your actual DB)
DB_PATH = "./data/app.db"

@app.list_tools()
async def handle_list_tools() -> list[types.Tool]:
    """List available tools."""
    return [
        types.Tool(
            name="query_database",
            description="Execute a read-only SQL query against the database",
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "SQL SELECT query to execute"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum rows to return (default: 100)",
                        "default": 100
                    }
                },
                "required": ["sql"]
            }
        ),
        types.Tool(
            name="list_tables",
            description="List all tables in the database with their schemas",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        )
    ]


@app.call_tool()
async def handle_call_tool(
    name: str,
    arguments: dict | None
) -> list[types.TextContent]:
    """Handle tool execution."""
    
    if name == "query_database":
        sql = arguments.get("sql", "")
        limit = arguments.get("limit", 100)
        
        # Safety: only allow SELECT
        if not sql.strip().upper().startswith("SELECT"):
            return [types.TextContent(
                type="text",
                text="Error: Only SELECT queries are allowed"
            )]
        
        # Add LIMIT if not present
        if "LIMIT" not in sql.upper():
            sql = f"{sql} LIMIT {limit}"
        
        try:
            conn = sqlite3.connect(DB_PATH)
            conn.row_factory = sqlite3.Row
            cursor = conn.execute(sql)
            rows = cursor.fetchall()
            conn.close()
            
            if not rows:
                return [types.TextContent(type="text", text="No results found")]
            
            # Format as markdown table
            headers = rows[0].keys()
            table = "| " + " | ".join(headers) + " |\n"
            table += "| " + " | ".join(["---"] * len(headers)) + " |\n"
            for row in rows:
                table += "| " + " | ".join(str(v) for v in row) + " |\n"
            
            return [types.TextContent(
                type="text",
                text=f"Query returned {len(rows)} rows:\n\n{table}"
            )]
            
        except Exception as e:
            return [types.TextContent(
                type="text", 
                text=f"Query error: {str(e)}"
            )]
    
    elif name == "list_tables":
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.execute(
            "SELECT name, sql FROM sqlite_master WHERE type='table'"
        )
        tables = cursor.fetchall()
        conn.close()
        
        result = ""
        for name, schema in tables:
            result += f"### {name}\n```sql\n{schema}\n```\n\n"
        
        return [types.TextContent(type="text", text=result or "No tables found")]
    
    raise ValueError(f"Unknown tool: {name}")


async def main():
    """Run the MCP server."""
    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="database-server",
                server_version="1.0.0",
                capabilities=app.get_capabilities(
                    notification_options=None,
                    experimental_capabilities={}
                )
            )
        )

if __name__ == "__main__":
    asyncio.run(main())

@app.list_resources()
async def handle_list_resources() -> list[types.Resource]:
    """Expose files as resources."""
    resources = []
    for f in Path("./docs").glob("*.md"):
        resources.append(types.Resource(
            uri=f"file://{f.absolute()}",
            name=f.name,
            mimeType="text/markdown"
        ))
    return resources

@app.read_resource()
async def handle_read_resource(uri: str) -> str:
    """Read a resource by URI."""
    path = uri.replace("file://", "")
    return Path(path).read_text()


@app.list_prompts()
async def handle_list_prompts() -> list[types.Prompt]:
    """Provide reusable prompt templates."""
    return [
        types.Prompt(
            name="analyze-schema",
            description="Analyze database schema and suggest improvements",
            arguments=[
                types.PromptArgument(
                    name="table_name",
                    description="Specific table to analyze (optional)",
                    required=False
                )
            ]
        )
    ]