#!/usr/bin/env python3
import json
import subprocess
from fastmcp import FastMCP

mcp = FastMCP('ClusterLogQuery')

@mcp.tool()
async def search_cluster_logs(keyword: str, unit: str = '*', limit: int = 500):
    """
    Search centralized systemd logs collected on this server.

    Args:
        keyword: The string to search for (e.g. 'oom-kill', 'failed').
        unit: The systemd unit to filter (e.g. 'sshd'). If '*', matches all.
        limit: The number of log lines to return (Max 1000).
    """
    
    limit = min(max(limit, 10), 1000)

    cmd = [
        'journalctl',
        '-n', str(limit),
        '-o', 'json',
        '--no-pager',
        '--directory', '/var/log/journal/remote'
    ]

    if unit != '*':
        cmd.append(f'UNIT={unit}')
    if keyword:
        cmd.append(keyword)

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, _ = proc.communicate()

        logs = []
        for line in stdout.decode('utf-8').strip().split('\n'):
            if not line: continue
            try:
                log = json.loads(line)
                logs.append({
                    'timestamp': log.get('TIME_ISO8601'),
                    'hostname': log.get('_HOSTNAME'),
                    'unit': log.get('_COMM'),
                    'message': log.get('MESSAGE')
                })
            except: pass

        return json.dumps(logs)

    except Exception as e:
        return f'Error querying logs: {str(e)}'

def main():
    mcp.run(transport='sse', host='0.0.0.0', port=9500)

if __name__ == '__main__':
    main()