server.py

"""HTTP API server + Raft node entry point.

Endpoints:
    PUT /key/<key>       - Set a key (body = value as plain text)
    GET /key/<key>       - Get a key's value
    DELETE /key/<key>    - Delete a key
    GET /status          - Node status (term, state, leader, log length)

Uses http.server.ThreadingHTTPServer for concurrent request handling.

When --http-port is not given, falls back to a simple stdin interface.
"""

import http.server
import json
import argparse
import threading
import logging
import sys

from raft.node import RaftNode
from raft.state_machine import KVStateMachine
from raft.persistence import Persister

logger = logging.getLogger(__name__)


class RaftHTTPHandler(http.server.BaseHTTPRequestHandler):
    """Request handler for the KV store HTTP API."""

    @property
    def raft_node(self) -> RaftNode:
        return self.server.raft_node

    @property
    def kv_store(self) -> KVStateMachine:
        return self.server.kv_store

    @property
    def peer_http_ports(self) -> dict:
        return self.server.peer_http_ports

    def do_PUT(self) -> None:
        key = self._parse_key()
        if key is None:
            self._send_json(400, {"error": "invalid path, use /key/<key>"})
            return

        content_length = int(self.headers.get("Content-Length", 0))
        value = self.rfile.read(content_length).decode("utf-8") if content_length else ""

        client_id = self.headers.get("X-Client-ID")
        request_id = self.headers.get("X-Request-ID")

        command = {"op": "set", "key": key, "value": value}
        success, result = self.raft_node.submit_command(
            command, client_id=client_id, request_id=request_id
        )

        if not success:
            if result and isinstance(result, str) and result != "timeout or leadership lost":
                # result is leader_id -- redirect
                self._redirect_to_leader(result)
            else:
                self._send_json(503, {"error": "not leader or timeout", "leader": result})
            return

        self._send_json(200, {"key": key, "value": result})

    def do_GET(self) -> None:
        path = self.path

        if path == "/status":
            status = self.raft_node.get_status()
            self._send_json(200, status)
            return

        key = self._parse_key()
        if key is None:
            self._send_json(400, {"error": "invalid path, use /key/<key> or /status"})
            return

        # Linearizable read: confirm leadership first
        if not self.raft_node.confirm_leadership():
            leader = self.raft_node.leader_id
            if leader:
                self._redirect_to_leader(leader)
            else:
                self._send_json(503, {"error": "no leader"})
            return

        value = self.kv_store.get(key)
        if value is None:
            self._send_json(404, {"error": "key not found"})
        else:
            self._send_json(200, {"key": key, "value": value})

    def do_DELETE(self) -> None:
        key = self._parse_key()
        if key is None:
            self._send_json(400, {"error": "invalid path, use /key/<key>"})
            return

        client_id = self.headers.get("X-Client-ID")
        request_id = self.headers.get("X-Request-ID")

        command = {"op": "delete", "key": key}
        success, result = self.raft_node.submit_command(
            command, client_id=client_id, request_id=request_id
        )

        if not success:
            if result and isinstance(result, str) and result != "timeout or leadership lost":
                self._redirect_to_leader(result)
            else:
                self._send_json(503, {"error": "not leader or timeout", "leader": result})
            return

        self._send_json(200, {"key": key, "deleted": result})

    def _parse_key(self) -> str | None:
        parts = self.path.strip("/").split("/", 1)
        if len(parts) == 2 and parts[0] == "key":
            return parts[1]
        return None

    def _redirect_to_leader(self, leader_id: str) -> None:
        http_port = self.peer_http_ports.get(leader_id)
        if http_port:
            location = f"http://127.0.0.1:{http_port}{self.path}"
            body = json.dumps({"leader_id": leader_id, "redirect": location}).encode("utf-8")
            self.send_response(307)
            self.send_header("Location", location)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        else:
            self._send_json(503, {"error": "leader unknown", "leader_id": leader_id})

    def _send_json(self, code: int, data: dict) -> None:
        body = json.dumps(data).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, format, *args):
        logger.debug(format, *args)


class RaftHTTPServer(http.server.ThreadingHTTPServer):
    def __init__(self, addr, handler_class, raft_node, kv_store, peer_http_ports):
        self.raft_node = raft_node
        self.kv_store = kv_store
        self.peer_http_ports = peer_http_ports
        super().__init__(addr, handler_class)


def stdin_command_loop(node: RaftNode, kv_store: KVStateMachine) -> None:
    """Simple stdin interface for testing without HTTP."""
    print("Raft node ready. Commands: set <key> <value>, get <key>, delete <key>, status, quit")
    while True:
        try:
            line = input("> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not line:
            continue
        parts = line.split(None, 2)
        cmd = parts[0].lower()

        if cmd == "quit":
            break
        elif cmd == "status":
            print(json.dumps(node.get_status(), indent=2))
        elif cmd == "set" and len(parts) >= 3:
            success, result = node.submit_command({"op": "set", "key": parts[1], "value": parts[2]})
            print(f"{'OK' if success else 'FAIL'}: {result}")
        elif cmd == "get" and len(parts) >= 2:
            value = kv_store.get(parts[1])
            print(f"Value: {value}")
        elif cmd == "delete" and len(parts) >= 2:
            success, result = node.submit_command({"op": "delete", "key": parts[1]})
            print(f"{'OK' if success else 'FAIL'}: {result}")
        else:
            print("Unknown command. Use: set <key> <value>, get <key>, delete <key>, status, quit")


def parse_peers(peers_str: str) -> dict[str, tuple[str, int]]:
    """Parse peer string like 'node2=127.0.0.1:5002,node3=127.0.0.1:5003'"""
    if not peers_str:
        return {}
    peers = {}
    for part in peers_str.split(","):
        name, addr = part.strip().split("=")
        host, port = addr.rsplit(":", 1)
        peers[name.strip()] = (host.strip(), int(port))
    return peers


def main():
    parser = argparse.ArgumentParser(description="Raft KV store node")
    parser.add_argument("--id", required=True, help="Node ID (e.g., node1)")
    parser.add_argument("--host", default="127.0.0.1", help="Bind address")
    parser.add_argument("--port", type=int, required=True, help="RPC port")
    parser.add_argument("--http-port", type=int, help="HTTP API port (omit for stdin mode)")
    parser.add_argument("--peers", default="", help="Peer list: id=host:port,...")
    parser.add_argument("--data-dir", help="Persistence directory")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format=f"%(asctime)s [{args.id}] %(levelname)s %(name)s: %(message)s",
    )

    peers = parse_peers(args.peers)

    kv_store = KVStateMachine()
    persister = Persister(args.data_dir) if args.data_dir else None

    node = RaftNode(
        node_id=args.id,
        host=args.host,
        port=args.port,
        peers=peers,
        apply_callback=kv_store.apply,
        snapshot_callback=kv_store.get_snapshot,
        load_snapshot_callback=kv_store.load_snapshot,
        persister=persister,
    )

    node.start()

    try:
        if args.http_port:
            # Build peer HTTP port mapping (convention: HTTP port = RPC port + 3000)
            peer_http_ports = {}
            for peer_id, (_, rpc_port) in peers.items():
                peer_http_ports[peer_id] = rpc_port + 3000

            http_server = RaftHTTPServer(
                (args.host, args.http_port),
                RaftHTTPHandler,
                node,
                kv_store,
                peer_http_ports,
            )
            logger.info("HTTP server on %s:%d", args.host, args.http_port)
            http_server.serve_forever()
        else:
            stdin_command_loop(node, kv_store)
    except KeyboardInterrupt:
        pass
    finally:
        node.stop()


if __name__ == "__main__":
    main()