#!/usr/bin/env python3
import argparse
import gzip
import json
import os
import re
import socket
import sys
import urllib.error
import urllib.parse
import urllib.request
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from pathlib import Path


CHINA_TZ = timezone(timedelta(hours=8))
DOMAINS = ("sc.jiansun.vip", "s3.jiansun.vip", "cp.jiansun.vip")
LOG_RE = re.compile(
    r'^(?P<remote>\S+) \S+ \S+ \[(?P<time>[^\]]+)\] '
    r'"(?P<method>[A-Z]+) (?P<target>\S+) HTTP/[^"]+" '
    r'(?P<status>\d{3}) (?P<body_bytes>\S+)(?P<rest>.*)$'
)


def parse_args():
    parser = argparse.ArgumentParser(description="Aggregate public traffic nginx logs and push rollups.")
    parser.add_argument("--log-dir", default="/var/log/nginx")
    parser.add_argument("--endpoint", default="http://127.0.0.1:3000/api/internal/public-traffic/rollups")
    parser.add_argument("--token-env", default="PUBLIC_TRAFFIC_ROLLUP_TOKEN")
    parser.add_argument("--granularity", choices=("5m", "day"), required=True)
    parser.add_argument("--start-date", help="China business date, YYYY-MM-DD")
    parser.add_argument("--end-date", help="China business date, YYYY-MM-DD")
    parser.add_argument("--source-server", default=socket.gethostname())
    parser.add_argument("--batch-size", type=int, default=1000)
    parser.add_argument("--domains", default=",".join(DOMAINS))
    return parser.parse_args()


def parse_date(value):
    return datetime.strptime(value, "%Y-%m-%d").date()


def default_dates(granularity):
    today = datetime.now(CHINA_TZ).date()
    if granularity == "5m":
        return today, today
    yesterday = today - timedelta(days=1)
    return yesterday, yesterday


def iter_log_paths(log_dir, domains):
    root = Path(log_dir)
    paths = []
    for domain in domains:
        paths.extend(root.glob(f"{domain}.access.log*"))
    return sorted(paths, key=lambda item: item.name)


def open_log(path):
    if str(path).endswith(".gz"):
        return gzip.open(path, "rt", encoding="utf-8", errors="replace")
    return open(path, "rt", encoding="utf-8", errors="replace")


def infer_domain(path):
    name = Path(path).name
    for domain in DOMAINS:
        if name.startswith(f"{domain}."):
            return domain
    return name.split(".access.log", 1)[0]


def normalize_target(target):
    try:
      parsed = urllib.parse.urlsplit(target)
      path = parsed.path or "/"
    except Exception:
      path = target.split("?", 1)[0] or "/"
    if not path.startswith("/"):
        path = "/" + path
    return path


def parse_size(value):
    if value == "-":
        return 0
    try:
        return max(0, int(value))
    except ValueError:
        return 0


def parse_extra_int(rest, key):
    match = re.search(rf"{re.escape(key)}=(\d+)", rest)
    return int(match.group(1)) if match else None


def status_bucket(status):
    if 200 <= status <= 299:
        return "2xx"
    if 300 <= status <= 399:
        return "3xx"
    if 400 <= status <= 499:
        return "4xx"
    if 500 <= status <= 599:
        return "5xx"
    return "other"


def parse_line(line, domain):
    match = LOG_RE.match(line)
    if not match:
        return None
    try:
        timestamp = datetime.strptime(match.group("time"), "%d/%b/%Y:%H:%M:%S %z").astimezone(CHINA_TZ)
        status = int(match.group("status"))
    except ValueError:
        return None

    rest = match.group("rest") or ""
    response_bytes = parse_extra_int(rest, "bytes_sent")
    if response_bytes is None:
        response_bytes = parse_size(match.group("body_bytes"))
    request_bytes = parse_extra_int(rest, "request_length")

    return {
        "timestamp": timestamp,
        "businessDate": timestamp.date().isoformat(),
        "domain": domain,
        "method": match.group("method").upper(),
        "path": normalize_target(match.group("target")),
        "statusBucket": status_bucket(status),
        "responseBytes": response_bytes,
        "requestBytes": request_bytes,
    }


def bucket_for(timestamp, granularity):
    if granularity == "day":
        start = timestamp.replace(hour=0, minute=0, second=0, microsecond=0)
        end = start + timedelta(days=1)
        return start, end
    minute = (timestamp.minute // 5) * 5
    start = timestamp.replace(minute=minute, second=0, microsecond=0)
    end = start + timedelta(minutes=5)
    return start, end


def iso_utc(value):
    return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")


def aggregate(args):
    domains = tuple(item.strip() for item in args.domains.split(",") if item.strip())
    start_default, end_default = default_dates(args.granularity)
    start_date = parse_date(args.start_date) if args.start_date else start_default
    end_date = parse_date(args.end_date) if args.end_date else end_default

    rows = {}
    scanned_lines = 0
    matched_lines = 0
    for path in iter_log_paths(args.log_dir, domains):
        domain = infer_domain(path)
        try:
            handle = open_log(path)
        except OSError as exc:
            print(f"WARN: cannot open {path}: {exc}", file=sys.stderr)
            continue
        with handle:
            for line in handle:
                scanned_lines += 1
                entry = parse_line(line, domain)
                if not entry:
                    continue
                entry_date = parse_date(entry["businessDate"])
                if entry_date < start_date or entry_date > end_date:
                    continue
                matched_lines += 1
                bucket_start, bucket_end = bucket_for(entry["timestamp"], args.granularity)
                key = (
                    iso_utc(bucket_start),
                    iso_utc(bucket_end),
                    entry["businessDate"],
                    args.granularity,
                    entry["domain"],
                    entry["method"],
                    entry["path"],
                )
                row = rows.setdefault(
                    key,
                    {
                        "bucketStart": key[0],
                        "bucketEnd": key[1],
                        "businessDate": key[2],
                        "granularity": key[3],
                        "domain": key[4],
                        "method": key[5],
                        "path": key[6],
                        "responseBytes": 0,
                        "requestBytes": None,
                        "requestBytesKnownCount": 0,
                        "totalRequests": 0,
                        "statusCounts": defaultdict(int),
                        "sourceServer": args.source_server,
                        "generatedAt": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
                        "frozen": args.granularity == "day",
                    },
                )
                row["responseBytes"] += entry["responseBytes"]
                row["totalRequests"] += 1
                row["statusCounts"][entry["statusBucket"]] += 1
                if entry["requestBytes"] is not None:
                    row["requestBytes"] = (row["requestBytes"] or 0) + entry["requestBytes"]
                    row["requestBytesKnownCount"] += 1

    normalized_rows = []
    for row in rows.values():
        row["statusCounts"] = dict(row["statusCounts"])
        normalized_rows.append(row)

    return normalized_rows, scanned_lines, matched_lines


def post_rows(endpoint, token, rows, batch_size):
    accepted = 0
    for index in range(0, len(rows), batch_size):
        chunk = rows[index:index + batch_size]
        body = json.dumps({"rows": chunk}, ensure_ascii=False).encode("utf-8")
        request = urllib.request.Request(
            endpoint,
            data=body,
            method="POST",
            headers={
                "Content-Type": "application/json",
                "X-Internal-Job-Token": token,
            },
        )
        try:
            with urllib.request.urlopen(request, timeout=240) as response:
                data = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"POST failed HTTP {exc.code}: {detail}") from exc
        accepted += int(data.get("accepted", 0))
    return accepted


def main():
    args = parse_args()
    token = os.environ.get(args.token_env, "").strip()
    if not token:
        raise SystemExit(f"Missing token env: {args.token_env}")

    rows, scanned_lines, matched_lines = aggregate(args)
    if rows:
        accepted = post_rows(args.endpoint, token, rows, max(1, args.batch_size))
    else:
        accepted = 0

    print(json.dumps({
        "ok": True,
        "granularity": args.granularity,
        "rows": len(rows),
        "accepted": accepted,
        "scannedLines": scanned_lines,
        "matchedLines": matched_lines,
    }, ensure_ascii=False))


if __name__ == "__main__":
    main()