#!/usr/bin/env python3
import argparse
import csv
import json
import os
import re
import subprocess
import sys
import time
from datetime import datetime, timezone


RESULT_RE = re.compile(
    r"^(?P<docker_ts>\S+)\s+RESULT Модель (?P<model_id>\d+) с типом (?P<model_type>.+?): "
    r"(?P<prediction>\S+) \(probability=(?P<probability>[-+]?\d+(?:\.\d+)?)\)\s*$"
)
FREQ_RE = re.compile(r"(\d{3,5})")


def parse_args():
    parser = argparse.ArgumentParser(
        description="Capture NN inference RESULT logs into CSV or JSONL until time or size limit."
    )
    parser.add_argument(
        "--output",
        required=True,
        help="Output file path (.csv or .jsonl recommended).",
    )
    parser.add_argument(
        "--format",
        choices=("csv", "jsonl"),
        default="csv",
        help="Output format.",
    )
    parser.add_argument(
        "--minutes",
        type=float,
        default=0.0,
        help="Stop after this many minutes. 0 means no time limit.",
    )
    parser.add_argument(
        "--max-bytes",
        type=int,
        default=3 * 1024 * 1024 * 1024,
        help="Stop when output file reaches this size in bytes. Default: 3 GiB.",
    )
    parser.add_argument(
        "--since",
        default=None,
        help="Optional docker logs --since value, e.g. 20m, 2h, 2026-05-04T12:00:00.",
    )
    parser.add_argument(
        "--tail",
        type=int,
        default=0,
        help="How many previous log lines to include before following. Default: 0.",
    )
    parser.add_argument(
        "--compose-file",
        default="deploy/docker/docker-compose.yml",
        help="Path to docker compose file.",
    )
    parser.add_argument(
        "--service",
        default="dronedetector-nn-server",
        help="Docker compose service name.",
    )
    parser.add_argument(
        "--follow",
        action="store_true",
        help="Follow logs live. By default the script captures a finite history snapshot.",
    )
    return parser.parse_args()


def extract_freq(model_type):
    matches = FREQ_RE.findall(model_type)
    if not matches:
        return ""
    known_freqs = {"433", "750", "868", "915", "1200", "1500", "2400", "3300", "4500", "5200", "5800"}
    for value in matches:
        if value in known_freqs:
            return value
    return matches[0]


def parse_docker_timestamp(value):
    try:
        return datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError:
        return datetime.now(timezone.utc)


def open_output(path, fmt):
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    fh = open(path, "a", encoding="utf-8", newline="")
    writer = None
    if fmt == "csv":
        writer = csv.writer(fh)
        if fh.tell() == 0:
            writer.writerow(
                [
                    "docker_timestamp",
                    "event_time_iso",
                    "event_time_epoch",
                    "freq",
                    "model_id",
                    "model_type",
                    "prediction",
                    "probability",
                ]
            )
            fh.flush()
    return fh, writer


def write_record(fh, writer, fmt, record):
    if fmt == "csv":
        writer.writerow(
            [
                record["docker_timestamp"],
                record["event_time_iso"],
                record["event_time_epoch"],
                record["freq"],
                record["model_id"],
                record["model_type"],
                record["prediction"],
                record["probability"],
            ]
        )
    else:
        fh.write(json.dumps(record, ensure_ascii=False) + "\n")
    fh.flush()


def build_command(args):
    cmd = [
        "docker",
        "compose",
        "-f",
        args.compose_file,
        "logs",
        "--timestamps",
        "--no-log-prefix",
        args.service,
    ]
    if args.since:
        cmd[5:5] = ["--since", args.since]
    if int(args.tail) > 0:
        cmd[-1:-1] = ["--tail", str(args.tail)]
    if args.follow:
        cmd.insert(-1, "-f")
    return cmd


def main():
    args = parse_args()
    deadline = time.time() + (args.minutes * 60.0) if args.minutes > 0 else None

    fh, writer = open_output(args.output, args.format)
    cmd = build_command(args)

    print("Running:", " ".join(cmd), file=sys.stderr)
    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding="utf-8",
        errors="replace",
        bufsize=1,
    )

    captured = 0
    try:
        assert proc.stdout is not None
        for line in proc.stdout:
            if deadline is not None and time.time() >= deadline:
                print("Stopping: time limit reached", file=sys.stderr)
                break

            match = RESULT_RE.match(line.rstrip("\n"))
            if not match:
                continue

            docker_dt = parse_docker_timestamp(match.group("docker_ts"))
            event_dt = docker_dt.astimezone()
            model_type = match.group("model_type")
            record = {
                "docker_timestamp": match.group("docker_ts"),
                "event_time_iso": event_dt.isoformat(timespec="seconds"),
                "event_time_epoch": round(docker_dt.timestamp(), 3),
                "freq": extract_freq(model_type),
                "model_id": int(match.group("model_id")),
                "model_type": model_type,
                "prediction": match.group("prediction"),
                "probability": float(match.group("probability")),
            }
            write_record(fh, writer, args.format, record)
            captured += 1

            if fh.tell() >= args.max_bytes:
                print("Stopping: file size limit reached", file=sys.stderr)
                break
    finally:
        try:
            proc.terminate()
        except Exception:
            pass
        try:
            proc.wait(timeout=5)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass
        fh.close()

    print(f"Captured {captured} inference results into {args.output}", file=sys.stderr)


if __name__ == "__main__":
    main()