#!/usr/bin/env python3
#
# @File: media-server.py
# @Date: 2026-05-21
#
# Local media server replicating production URL-based background removal.
#
# - GET /dealer_01.jpg              → serves original image
# - GET /dealer_01.jpg?bgRemove=true → removes background, returns PNG (cached)
#
# Usage:
#   python3 media-server.py --media-root /tmp/media_local --port 8080
#
# Then open in browser:
#   http://localhost:8080/dealer_01.jpg
#   http://localhost:8080/dealer_01.jpg?bgRemove=true

import sys
import argparse
from pathlib import Path
from urllib.parse import urlparse, parse_qs
from http.server import HTTPServer, BaseHTTPRequestHandler

try:
    from rembg import remove, new_session
except ImportError:
    print("ERROR: Run:  pip install 'rembg[cpu]'")
    sys.exit(1)

MEDIA_ROOT = None
CACHE_DIR  = None
SESSION    = None

MIME = {
    '.jpg':  'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.png':  'image/png',
    '.webp': 'image/webp',
}


def load_model():
    global SESSION
    print("Loading BiRefNet model (CPU) ...")
    SESSION = new_session('birefnet-general', providers=['CPUExecutionProvider'])
    print("Model ready. Server starting...\n")


class MediaHandler(BaseHTTPRequestHandler):

    def log_message(self, fmt, *args):
        print(f"  {self.address_string()} → {fmt % args}")

    def do_GET(self):
        parsed   = urlparse(self.path)
        params   = parse_qs(parsed.query)
        rel_path = parsed.path.lstrip('/')

        img_path = MEDIA_ROOT / rel_path
        if not img_path.exists() or img_path.suffix.lower() not in MIME:
            self.send_error(404, f"Not found: {rel_path}")
            return

        bg_remove = params.get('bgRemove', ['false'])[0].lower() == 'true'

        if not bg_remove:
            # Serve original
            data = img_path.read_bytes()
            self.send_response(200)
            self.send_header('Content-Type', MIME[img_path.suffix.lower()])
            self.send_header('Content-Length', str(len(data)))
            self.end_headers()
            self.wfile.write(data)
            return

        # Background removal — check cache first
        cache_path = CACHE_DIR / (img_path.stem + '_nobg.png')
        if not cache_path.exists():
            print(f"  Processing {img_path.name} with BiRefNet ...")
            result = remove(img_path.read_bytes(), session=SESSION)
            cache_path.write_bytes(result)
            print(f"  Cached → {cache_path.name}")

        data = cache_path.read_bytes()
        self.send_response(200)
        self.send_header('Content-Type', 'image/png')
        self.send_header('Content-Length', str(len(data)))
        self.end_headers()
        self.wfile.write(data)


def main():
    global MEDIA_ROOT, CACHE_DIR

    parser = argparse.ArgumentParser(description='Local media server with on-demand background removal')
    parser.add_argument('--media-root', default='/tmp/media_local', help='Folder containing images')
    parser.add_argument('--cache-dir',  default='',                 help='Folder for cached bg-removed PNGs (default: media-root/cache)')
    parser.add_argument('--port',       type=int, default=8080,     help='Port to listen on (default: 8080)')
    args = parser.parse_args()

    MEDIA_ROOT = Path(args.media_root)
    CACHE_DIR  = Path(args.cache_dir) if args.cache_dir else MEDIA_ROOT / 'cache'
    CACHE_DIR.mkdir(parents=True, exist_ok=True)

    print("=" * 50)
    print("  Local Media Server")
    print("=" * 50)
    print(f"  Media root : {MEDIA_ROOT}")
    print(f"  Cache dir  : {CACHE_DIR}")
    print(f"  Port       : {args.port}")
    print()
    print(f"  Original : http://localhost:{args.port}/dealer_01.jpg")
    print(f"  No BG    : http://localhost:{args.port}/dealer_01.jpg?bgRemove=true")
    print()

    load_model()

    server = HTTPServer(('0.0.0.0', args.port), MediaHandler)
    print(f"Listening on http://localhost:{args.port} — Ctrl+C to stop\n")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.")


if __name__ == '__main__':
    main()