#!/usr/bin/env python3
import argparse
import json
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse, parse_qs


def parse_phases(spec: str):
    phases = []
    total = 0
    for part in spec.split(','):
        part = part.strip()
        if not part:
            continue
        name, seconds = part.split(':', 1)
        dur = int(seconds)
        if dur <= 0:
            continue
        total += dur
        phases.append((name.strip(), total))
    if not phases:
        phases = [("data", 3600)]
    return phases


def current_phase(start_ts: float, phases):
    elapsed = int(time.time() - start_ts)
    for name, end_at in phases:
        if elapsed < end_at:
            return name, elapsed
    return phases[-1][0], elapsed


class Handler(BaseHTTPRequestHandler):
    server_version = "mock-o11y/0.1"

    def _json(self, code: int, payload: dict):
        raw = json.dumps(payload).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def _text(self, code: int, payload: str):
        raw = payload.encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "text/plain")
        self.send_header("Content-Length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def log_message(self, fmt, *args):
        return

    def do_GET(self):
        parsed = urlparse(self.path)
        qs = parse_qs(parsed.query)
        phase, elapsed = current_phase(self.server.start_ts, self.server.phases)

        if parsed.path in ("/healthz", "/health"):
            self._json(200, {"ok": True, "phase": phase, "elapsed": elapsed})
            return

        if parsed.path.startswith("/api/v1/clusters"):
            host = self.server.internal_read_host
            uri = f"http://{host}/internal/metrics"
            self._json(200, {
                "code": 0,
                "message": "",
                "data": {
                    "clusters": [
                        {
                            "id": str(self.server.cluster_id),
                            "vendor": self.server.vendor,
                            "region": self.server.region,
                            "tenantID": str(self.server.tenant_id),
                            "projectID": str(self.server.project_id),
                            "applicationID": str(self.server.cluster_id),
                            "bizType": self.server.biz_type,
                            "internalReadUri": uri,
                            "name": "mock-cluster",
                            "labels": [],
                            "version": "mock",
                            "shareMode": "mock",
                            "note": "mock",
                        }
                    ],
                    "continueToken": "",
                    "remainingItemCount": 0,
                },
            })
            return

        if parsed.path.endswith("/query_range") or parsed.path.endswith("/query"):
            promql = qs.get("query", [""])[0]
            now = int(time.time())

            # Keep keyspace lookup always available so Essential/Starter RU queries
            # can resolve resource_group even during empty/error phases.
            if parsed.path.endswith("/query") and (
                "topk(1" in promql
                and (
                    "tidb_server_connections" in promql
                    or "tidb_server_tokens" in promql
                    or "tidb_server_handle_query_duration_seconds_sum" in promql
                )
            ):
                self._json(200, {
                    "status": "success",
                    "data": {
                        "resultType": "vector",
                        "result": [
                            {
                                "metric": {
                                    "__name__": "mock_keyspace_lookup",
                                    "cluster_id": str(self.server.cluster_id),
                                    "keyspace_id": str(self.server.keyspace_id),
                                    "chaos_run": self.server.run_id,
                                },
                                "value": [now, "1"],
                            }
                        ],
                    },
                })
                return

            if phase == "error":
                self._text(500, f"mock upstream error phase={phase} elapsed={elapsed}")
                return

            if phase == "empty":
                self._json(200, {
                    "status": "success",
                    "data": {
                        "resultType": "matrix",
                        "result": [],
                    },
                })
                return

            if parsed.path.endswith("/query"):
                self._json(200, {
                    "status": "success",
                    "data": {
                        "resultType": "vector",
                        "result": [
                            {
                                "metric": {
                                    "__name__": "mock_metric",
                                    "cluster_id": str(self.server.cluster_id),
                                    "chaos_run": self.server.run_id,
                                },
                                "value": [now, "1"],
                            }
                        ],
                    },
                })
                return

            start = int(float(qs.get("start", [str(now - 60)])[0]))
            end = int(float(qs.get("end", [str(now)])[0]))
            step = 15
            vals = []
            t = start
            while t <= end:
                vals.append([t, "1"])
                t += step

            self._json(200, {
                "status": "success",
                "data": {
                    "resultType": "matrix",
                    "result": [
                        {
                            "metric": {
                                "__name__": "mock_metric",
                                "cluster_id": str(self.server.cluster_id),
                                "chaos_run": self.server.run_id,
                            },
                            "values": vals,
                        }
                    ],
                },
            })
            return

        self._text(404, f"not found: {parsed.path}\n")


def main():
    ap = argparse.ArgumentParser(description="Mock O11Y global/regional server for local chaos replay")
    ap.add_argument("--port", type=int, default=19090)
    ap.add_argument("--phases", default="empty:45,error:45,data:90")
    ap.add_argument("--cluster-id", required=True)
    ap.add_argument("--tenant-id", required=True)
    ap.add_argument("--project-id", default="0")
    ap.add_argument("--vendor", default="aws")
    ap.add_argument("--region", default="us-east-1")
    ap.add_argument("--biz-type", default="dev-tidbcloud")
    ap.add_argument("--run-id", default="local-chaos")
    ap.add_argument("--keyspace-id", default="1")
    args = ap.parse_args()

    server = ThreadingHTTPServer(("127.0.0.1", args.port), Handler)
    server.start_ts = time.time()
    server.phases = parse_phases(args.phases)
    server.cluster_id = args.cluster_id
    server.tenant_id = args.tenant_id
    server.project_id = args.project_id
    server.vendor = args.vendor
    server.region = args.region
    server.biz_type = args.biz_type
    server.run_id = args.run_id
    server.keyspace_id = args.keyspace_id
    server.internal_read_host = "www.ds.us-east-1-serverless.aws.observability.tidbcloud.com"

    print(json.dumps({
        "msg": "mock server started",
        "port": args.port,
        "phases": args.phases,
        "cluster_id": args.cluster_id,
        "run_id": args.run_id,
    }), flush=True)

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()