#!/usr/bin/env python3
"""
TCP log server — receives JSON-line logs from index-service and search-service.

Usage:
    python3 server.py                  # listen on 0.0.0.0:9999
    python3 server.py --port 9999
    python3 server.py --save logs.jsonl   # also save to file
"""

import argparse
import json
import logging
import socketserver
import sys
import threading
from datetime import datetime

COLORS = {
    "DEBUG":    "\033[36m",
    "INFO":     "\033[0m",
    "WARNING":  "\033[33m",
    "ERROR":    "\033[31m",
    "CRITICAL": "\033[35m",
}
RESET = "\033[0m"
SERVICE_COLOR = {
    "index-service":  "\033[34m",  # blue
    "search-service": "\033[32m",  # green
}

_save_file = None
_save_lock = threading.Lock()


def _format(entry: dict) -> str:
    ts = entry.get("ts", "")[:23].replace("T", " ")
    level = entry.get("level", "INFO")
    service = entry.get("service", "?")
    msg = entry.get("msg", "")

    lc = COLORS.get(level, "")
    sc = SERVICE_COLOR.get(service, "\033[0m")
    return f"{ts}  {sc}{service:<15}{RESET}  {lc}{level:<8}{RESET}  {msg}"


def _handle_line(raw: str) -> None:
    raw = raw.strip()
    if not raw:
        return
    try:
        entry = json.loads(raw)
    except json.JSONDecodeError:
        entry = {"ts": datetime.utcnow().isoformat(), "level": "INFO", "service": "?", "msg": raw}

    print(_format(entry), flush=True)

    if _save_file:
        with _save_lock:
            _save_file.write(raw + "\n")
            _save_file.flush()


class _Handler(socketserver.StreamRequestHandler):
    def handle(self) -> None:
        addr = self.client_address[0]
        print(f"\033[90m[+] connected: {addr}{RESET}", flush=True)
        try:
            for raw_bytes in self.rfile:
                try:
                    _handle_line(raw_bytes.decode("utf-8", errors="replace"))
                except Exception:
                    pass
        except Exception:
            pass
        print(f"\033[90m[-] disconnected: {addr}{RESET}", flush=True)


def main() -> None:
    global _save_file

    parser = argparse.ArgumentParser(description="TCP JSON-line log receiver")
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=9999)
    parser.add_argument("--save", metavar="FILE", help="Also save raw JSON lines to this file")
    args = parser.parse_args()

    if args.save:
        _save_file = open(args.save, "a", encoding="utf-8")
        print(f"Saving logs to {args.save}", flush=True)

    server = socketserver.ThreadingTCPServer((args.host, args.port), _Handler)
    server.allow_reuse_address = True

    print(f"Listening on {args.host}:{args.port} ...\n", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.")
    finally:
        server.server_close()
        if _save_file:
            _save_file.close()


if __name__ == "__main__":
    main()