import json
import os
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlsplit

import pymysql


PORT = int(os.environ.get("PORT", "8080"))
DB_CONFIG = {
    "host": os.environ.get("DB_HOST", "127.0.0.1"),
    "port": int(os.environ.get("DB_PORT", "3306")),
    "database": os.environ.get("DB_NAME", "maria"),
    "user": os.environ.get("DB_USER", ""),
    "password": os.environ.get("DB_PASSWORD", ""),
    "charset": "utf8mb4",
    "connect_timeout": 3,
    "read_timeout": 3,
    "write_timeout": 3,
    "autocommit": True,
}


def database_ready():
    connection = None
    cursor = None
    try:
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        cursor.execute("SELECT 1")
        return cursor.fetchone() == (1,)
    except (pymysql.MySQLError, OSError):
        return False
    finally:
        if cursor is not None:
            cursor.close()
        if connection is not None:
            connection.close()


class Handler(BaseHTTPRequestHandler):
    server_version = "MariaDBCheck/1.0"
    sys_version = ""

    def send_json(self, status, payload):
        body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        path = urlsplit(self.path).path
        if path == "/healthz":
            self.send_json(HTTPStatus.OK, {"status": "ok"})
            return

        if path in ("/", "/readyz"):
            ready = database_ready()
            status = HTTPStatus.OK if ready else HTTPStatus.SERVICE_UNAVAILABLE
            self.send_json(
                status,
                {
                    "service": "maria-http",
                    "status": "ready" if ready else "unavailable",
                    "database": "connected" if ready else "unavailable",
                },
            )
            return

        self.send_json(HTTPStatus.NOT_FOUND, {"error": "not_found"})


if __name__ == "__main__":
    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
    server.serve_forever()