#!/bin/python
from http.server import HTTPServer, SimpleHTTPRequestHandler
from http.client import HTTPConnection
from urllib.parse import urlparse, parse_qs
from functools import partial
import os


# This kinda mimics the behavior of my nginx config.
# Do not use it in prod!

LOCALHOST = "127.0.0.1"
SEARCH_SERVER_PORT = 8081
WEB_SERVER_PORT = 8000
ROOT_DIR = "www/"


class RewriteHandler(SimpleHTTPRequestHandler):
    def do_GET(self):
        parsed = urlparse(self.path)
        query = parse_qs(parsed.query)
        path = parsed.path

        if 'q' in query:
            self.forward_to_search_server(self.path)
            return
        elif 'ts' in query:
            ts_value = query['ts'][0]
            new_path = f"posts/{ts_value}.html"
            self.path = new_path
        elif parsed.path == "/rss":
            self.path = "/ptrace.dev.rss"
        else:
            file_path = self.translate_path(path)
            if not os.path.exists(file_path):
                file_path_html = self.translate_path(path + ".html")
                if os.path.exists(file_path_html):
                    self.path = path + ".html"

        return super().do_GET()

    def handle(self):
        try:
            super().handle()
        except BrokenPipeError:
            pass

    def send_error(self, code, message=None, explain=None):
        try:
            super().send_error(code, message, explain)
        except BrokenPipeError:
            pass

    def forward_to_search_server(self, path):
        try:
            conn = HTTPConnection(LOCALHOST, SEARCH_SERVER_PORT)
            conn.request("GET", path, headers={"Host": f"{LOCALHOST}:8000"})
            resp = conn.getresponse()

            self.send_response(resp.status)
            for key, value in resp.getheaders():
                if key.lower() not in ("transfer-encoding", "connection"):
                    self.send_header(key, value)
            self.end_headers()

            self.wfile.write(resp.read())
        except ConnectionRefusedError:
            self.send_error(502, "Search server unavailable")
        finally:
            conn.close()


if __name__ == "__main__":
    handler = partial(RewriteHandler, directory=ROOT_DIR)
    server = HTTPServer((LOCALHOST, WEB_SERVER_PORT), handler)

    try:
        print(f"Serving {ROOT_DIR} on http://localhost:{WEB_SERVER_PORT}")
        print("DO NOT USE IT IN PROD!")
        print("DO NOT USE IT IN PROD!")
        print("DO NOT USE IT IN PROD!")
        print("----------------------")
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopping server...")
        server.shutdown()
        server.server_close()