# db_server.py
import asyncio
import os
import json
from contextlib import asynccontextmanager
import asyncpg
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp import types

DB_URL = os.environ["DATABASE_URL"]
ALLOWED_SCHEMAS = {"public", "analytics"}  # Security: restrict schema access

app = Server("database-tools")
db_pool: asyncpg.Pool | None = None

@asynccontextmanager
async def lifespan():
    global db_pool
    db_pool = await asyncpg.create_pool(DB_URL, min_size=1, max_size=5)
    yield
    await db_pool.close()

@app.list_tools()
async def list_tools() -> list[types.Tool]:
    return [
        types.Tool(
            name="query_database",
            description=(
                "Run a read-only SQL query against the database. "
                "Only SELECT statements are allowed. "
                "Returns results as JSON."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "SQL SELECT query to execute"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum rows to return (default 100, max 1000)",
                        "default": 100
                    }
                },
                "required": ["sql"]
            }
        ),
        types.Tool(
            name="list_tables",
            description="List all accessible tables and their column definitions",
            inputSchema={
                "type": "object",
                "properties": {
                    "schema": {
                        "type": "string",
                        "description": "Schema name (default: public)",
                        "default": "public"
                    }
                }
            }
        )
    ]

@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[types.TextContent]:
    if name == "query_database":
        return await handle_query(arguments)
    elif name == "list_tables":
        return await handle_list_tables(arguments)
    raise ValueError(f"Unknown tool: {name}")

async def handle_query(args: dict) -> list[types.TextContent]:
    sql = args["sql"].strip()
    limit = min(int(args.get("limit", 100)), 1000)

    # Security: only allow SELECT
    if not sql.upper().startswith("SELECT"):
        return [types.TextContent(
            type="text",
            text="Error: Only SELECT queries are permitted"
        )]

    # Security: block schema access outside allowed set
    for schema in ALLOWED_SCHEMAS:
        if schema not in sql.lower() and "." in sql:
            pass  # Will rely on DB permissions for enforcement

    try:
        async with db_pool.acquire() as conn:
            # Run in read-only transaction
            async with conn.transaction(readonly=True):
                rows = await conn.fetch(
                    f"SELECT * FROM ({sql}) AS q LIMIT $1",
                    limit
                )

        result = [dict(row) for row in rows]
        return [types.TextContent(
            type="text",
            text=json.dumps({
                "rows": result,
                "count": len(result),
                "truncated": len(result) == limit
            }, indent=2, default=str)
        )]
    except asyncpg.PostgresError as e:
        return [types.TextContent(type="text", text=f"Database error: {e}")]

async def handle_list_tables(args: dict) -> list[types.TextContent]:
    schema = args.get("schema", "public")
    if schema not in ALLOWED_SCHEMAS:
        return [types.TextContent(type="text", text=f"Schema '{schema}' not accessible")]

    async with db_pool.acquire() as conn:
        tables = await conn.fetch("""
            SELECT table_name,
                   json_agg(json_build_object(
                       'column', column_name,
                       'type', data_type,
                       'nullable', is_nullable
                   ) ORDER BY ordinal_position) AS columns
            FROM information_schema.columns
            WHERE table_schema = $1
            GROUP BY table_name
            ORDER BY table_name
        """, schema)

    result = {row["table_name"]: json.loads(row["columns"]) for row in tables}
    return [types.TextContent(type="text", text=json.dumps(result, indent=2))]

async def main():
    async with lifespan():
        async with stdio_server() as (read_stream, write_stream):
            await app.run(read_stream, write_stream, app.create_initialization_options())

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

# log_resource_server.py
import asyncio
import os
from pathlib import Path
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp import types

app = Server("log-resources")
LOG_DIR = Path(os.environ.get("LOG_DIR", "/var/log/app"))

@app.list_resources()
async def list_resources() -> list[types.Resource]:
    resources = []
    for log_file in LOG_DIR.glob("*.log"):
        resources.append(types.Resource(
            uri=f"file://{log_file}",
            name=log_file.name,
            description=f"Log file: {log_file.name}",
            mimeType="text/plain"
        ))
    return resources

@app.read_resource()
async def read_resource(uri: str) -> str:
    path = Path(uri.replace("file://", ""))

    # Security: ensure path is within LOG_DIR
    try:
        path.resolve().relative_to(LOG_DIR.resolve())
    except ValueError:
        raise ValueError(f"Access denied: {uri}")

    if not path.exists():
        raise FileNotFoundError(f"Log file not found: {path.name}")

    # Return last 200 lines
    lines = path.read_text().splitlines()
    return "\n".join(lines[-200:])

async def main():
    async with stdio_server() as (read_stream, write_stream):
        await app.run(read_stream, write_stream, app.create_initialization_options())

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