#!/usr/bin/env python3
"""
Reverse proxy on port 8788.
- /yt-dlp-web/*    -> 127.0.0.1:48711 (yt-dlp-web FastAPI server)
- everything else   -> serve from /opt/data/diary/ as static files

Replaces both the old 8788->8789 forwarder AND the 8789 diary_server.
"""
import http.server
import mimetypes
import os
import socket
import socketserver
import sys
import time


LISTEN = ("0.0.0.0", 8788)
SERVE_DIR = "/opt/data/diary"
YT_DLP_WEB_HOST = "127.0.0.1"
YT_DLP_WEB_PORT = 48711
DIARY_SERVER_HOST = "127.0.0.1"
DIARY_SERVER_PORT = 8789
LOG = "/opt/data/logs/diary_port_forward.log"


def is_yt_dlp_web_path(path: str) -> bool:
    return path.startswith("/yt-dlp-web")


def peek_request_line(client, max_bytes=4096):
    """Read the HTTP request line + buffered data so we can route
    AND replay the bytes to the upstream.
    """
    try:
        client.settimeout(1.0)
        data = b""
        while b"\r\n" not in data and len(data) < max_bytes:
            try:
                chunk = client.recv(1024)
            except socket.timeout:
                break
            if not chunk:
                break
            data += chunk
        if not data:
            return None, None, b""
        first_line_end = data.find(b"\r\n")
        if first_line_end < 0:
            return None, None, data
        first_line = data[:first_line_end].decode("latin-1", errors="ignore")
        parts = first_line.split()
        if len(parts) < 2:
            return None, None, data
        method = parts[0]
        path = parts[1]
        return method, path, data
    except Exception:
        return None, None, b""


def rewrite_location_header(line: bytes, prefix: str) -> bytes:
    """Inject `prefix` into a Location header (absolute URL or relative path)."""
    if not line.lower().startswith(b"location:"):
        return line
    loc = line.split(b":", 1)[1].strip()
    prefix_b = prefix.encode()
    if loc.startswith(b"/") and not loc.startswith(prefix_b + b"/") and loc != prefix_b:
        return b"Location: " + prefix_b + loc
    if loc.startswith(b"http://") or loc.startswith(b"https://"):
        try:
            scheme, rest = loc.split(b"://", 1)
            _, path = rest.split(b"/", 1)
            path = b"/" + path
            if not path.startswith(prefix_b + b"/") and path != prefix_b:
                new_path = prefix_b + path
                new_loc = scheme + b"://" + rest.split(b"/", 1)[0] + new_path
                return b"Location: " + new_loc
        except ValueError:
            pass
    return line


def proxy_to_yt_dlp_web(client, method, path, peeked):
    """Forward the request to the FastAPI server on 48711."""
    # Strip /yt-dlp-web prefix
    if path.startswith("/yt-dlp-web"):
        new_path = path[len("/yt-dlp-web"):]
        if not new_path:
            new_path = "/"
    else:
        new_path = path

    # Rewrite first line in peeked buffer
    if peeked and path != new_path:
        first_line_end = peeked.find(b"\r\n")
        if first_line_end > 0:
            old_first_line = peeked[:first_line_end]
            space_idx = old_first_line.find(b" ")
            if space_idx > 0:
                rest = old_first_line[space_idx + 1:]
                end_of_path = rest.find(b" ")
                if end_of_path > 0:
                    new_first_line = method.encode() + b" " + new_path.encode() + rest[end_of_path:]
                    peeked = new_first_line + peeked[first_line_end:]

    # Connect to upstream
    try:
        upstream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        upstream.settimeout(60)
        upstream.connect((YT_DLP_WEB_HOST, YT_DLP_WEB_PORT))
    except OSError as e:
        try:
            client.send_error(502, f"yt-dlp-web upstream error: {e}")
        except Exception:
            pass
        client.close()
        return

    # Send replay buffer
    if peeked:
        try:
            upstream.sendall(peeked)
        except OSError:
            pass

    # Forward upstream -> client with Location header rewriting
    buf = b""
    try:
        while True:
            data = upstream.recv(4096)
            if not data:
                break
            buf += data
            while b"\r\n" in buf:
                line, _, rest = buf.partition(b"\r\n")
                if line.lower().startswith(b"location:"):
                    line = rewrite_location_header(line, "/yt-dlp-web")
                try:
                    client.wfile.write(line + b"\r\n")
                except OSError:
                    return
                buf = rest
    except OSError:
        pass
    finally:
        if buf:
            try:
                client.wfile.write(buf)
            except OSError:
                pass
        try:
            client.wfile.flush()
        except OSError:
            pass
        upstream.close()
        try:
            client.connection.close()
        except Exception:
            pass


def proxy_to_diary_server(client, method, path, peeked):
    """Forward /api/tomica/* to diary_server.py on 8789 (which has the real handlers)."""
    try:
        upstream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        upstream.settimeout(60)
        upstream.connect((DIARY_SERVER_HOST, DIARY_SERVER_PORT))
    except OSError as e:
        try:
            client.send_error(502, f"diary_server upstream error: {e}")
        except Exception:
            pass
        client.close()
        return

    # Send replay buffer (path stays the same — diary_server routes by self.path)
    if peeked:
        try:
            upstream.sendall(peeked)
        except OSError:
            pass

    # Pipe upstream -> client
    buf = b""
    try:
        while True:
            data = upstream.recv(4096)
            if not data:
                break
            buf += data
            while b"\r\n" in buf:
                line, _, rest = buf.partition(b"\r\n")
                try:
                    client.wfile.write(line + b"\r\n")
                except OSError:
                    return
                buf = rest
    except OSError:
        pass
    finally:
        if buf:
            try:
                client.wfile.write(buf)
            except OSError:
                pass
        try:
            client.wfile.flush()
        except OSError:
            pass
        upstream.close()
        try:
            client.connection.close()
        except Exception:
            pass


def serve_static_file(client, method, path, peeked):
    """Serve a static file from SERVE_DIR using mimetypes."""
    if path == "/" or path == "":
        path = "/diary_blog.html"
    if "?" in path:
        path = path.split("?", 1)[0]
    rel = path.lstrip("/")
    fs_path = os.path.join(SERVE_DIR, rel)
    if not os.path.exists(fs_path) or os.path.isdir(fs_path):
        client.send_error(404, f"Not found: {path}")
        return
    ctype, _ = mimetypes.guess_type(fs_path)
    if ctype is None:
        ctype = "application/octet-stream"
    try:
        with open(fs_path, "rb") as f:
            data = f.read()
    except OSError as e:
        client.send_error(500, str(e))
        return
    client.send_response(200)
    client.send_header("Content-Type", ctype)
    client.send_header("Content-Length", str(len(data)))
    client.send_header("Cache-Control", "no-cache")
    client.end_headers()
    client.wfile.write(data)


class Handler(http.server.BaseHTTPRequestHandler):
    """Routes:
       - /yt-dlp-web/*  -> 127.0.0.1:48711
       - everything else -> /opt/data/diary/ static file
    """

    def log_message(self, format, *args):
        # Silent
        return

    def parse_request(self):
        """Override: don't auto-parse. We do our own peek+route."""
        return False

    def handle(self):
        try:
            method, path, peeked = peek_request_line(self.connection)
            if not method:
                self.close_connection = True
                return
            # Populate BaseHTTPRequestHandler attributes expected by send_error/send_response
            self.command = method
            self.path = path
            first_line_end = peeked.find(b"\r\n")
            if first_line_end > 0:
                self.requestline = peeked[:first_line_end].decode("latin-1", errors="ignore")
            else:
                self.requestline = f"{method} {path} HTTP/1.1"
            # Best-effort version parse
            try:
                _verb, _p, ver = self.requestline.split(" ", 2)
                self.request_version = ver.split("/")[-1]
            except Exception:
                self.request_version = "1.1"

            # Apply / -> /diary_blog.html rewrite
            if path in ("/", "/diary"):
                path = "/diary_blog.html"

            # Inject /yt-dlp-web prefix for ALL yt-dlp-web routes (including
            # requests from internal JS that use /static/* without
            # the /yt-dlp-web/ prefix). This is more forgiving than exact-match.
            # /api/* is handled differently — see below.
            if path.startswith("/api/"):
                # EXCEPTION: /api/tomica/* is served by diary_server.py (port 8789)
                if path.startswith("/api/tomica/"):
                    proxy_to_diary_server(self, method, path, peeked)
                    return
                # EXCEPTION: /api/garage/* (Tomica v2) also served by diary_server
                if path.startswith("/api/garage/"):
                    proxy_to_diary_server(self, method, path, peeked)
                    return
                # EXCEPTION: /api/gba/* — diary_server local file handler
                if path.startswith("/api/gba/"):
                    proxy_to_diary_server(self, method, path, peeked)
                    return
                # All other /api/* — proxy to diary_server (port 8789) so its
                # own routing decides: image_gen, bureau, ledger, etc.
                # Without this, Hermes service worker fetches like /api/auth/status
                # get rewritten and break the chain.
                proxy_to_diary_server(self, method, path, peeked)
                return

            if path.startswith("/static/"):
                # EXCEPTION: /ledger/static/* — strip /ledger, proxy to diary_server
                # (which forwards to the Victoria Ledger Flask on 18790).
                if path.startswith("/ledger/static/"):
                    proxy_to_diary_server(self, method, path[len("/ledger"):], peeked)
                    return
                # Rewrite /static/* to /yt-dlp-web/static/* for the yt-dlp-web SPA
                new_path = "/yt-dlp-web" + path
                # Rewrite first line in peeked buffer
                first_line_end = peeked.find(b"\r\n")
                if first_line_end > 0:
                    old_first_line = peeked[:first_line_end]
                    space_idx = old_first_line.find(b" ")
                    if space_idx > 0:
                        rest = old_first_line[space_idx + 1:]
                        end_of_path = rest.find(b" ")
                        if end_of_path > 0:
                            new_first_line = method.encode() + b" " + new_path.encode() + rest[end_of_path:]
                            peeked = new_first_line + peeked[first_line_end:]
                path = new_path

            # 維港帳房 (Victoria Ledger) — reverse-proxy /ledger* to diary_server.py
            # (which forwards to the Flask sub-process on 18790).
            # Match /ledger, /ledger/, /ledger.html, /ledger?...
            if path == "/ledger" or path == "/ledger.html" or path.startswith("/ledger/") or path.startswith("/ledger?") or path.startswith("/ledger."):
                proxy_to_diary_server(self, method, path, peeked)
                return

            if is_yt_dlp_web_path(path):
                proxy_to_yt_dlp_web(self, method, path, peeked)
            else:
                serve_static_file(self, method, path, peeked)
        except Exception as e:
            print(f"[handler] error: {e}", file=sys.stderr, flush=True)
        finally:
            try:
                self.close_connection = True
                self.connection.close()
            except Exception:
                pass


class ThreadedHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True
    allow_reuse_address = True


def main():
    os.makedirs(os.path.dirname(LOG), exist_ok=True)
    server = ThreadedHTTPServer(LISTEN, Handler)
    msg = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] diary+yt-dlp-web server up on {LISTEN[0]}:{LISTEN[1]} (pid={os.getpid()})\n"
    print(msg, flush=True)
    with open(LOG, "a") as logfd:
        logfd.write(msg)
        logfd.flush()
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()