import json
import os
import subprocess
import sys
import time


def run(argv, timeout):
    started = time.monotonic()
    try:
        completed = subprocess.run(argv, capture_output=True, text=True, timeout=timeout, check=False)
        return {
            "exitCode": completed.returncode,
            "stdout": completed.stdout,
            "stderr": completed.stderr,
            "elapsedMs": round((time.monotonic() - started) * 1000),
            "timedOut": False,
        }
    except subprocess.TimeoutExpired as error:
        stdout = error.stdout.decode("utf-8", errors="replace") if isinstance(error.stdout, bytes) else (error.stdout or "")
        stderr = error.stderr.decode("utf-8", errors="replace") if isinstance(error.stderr, bytes) else (error.stderr or "")
        return {
            "exitCode": 124,
            "stdout": stdout,
            "stderr": stderr,
            "elapsedMs": round((time.monotonic() - started) * 1000),
            "timedOut": True,
        }


def json_object(text):
    lines = [line.strip() for line in text.splitlines() if line.strip().startswith("{")]
    for line in reversed(lines):
        try:
            value = json.loads(line)
            if isinstance(value, dict):
                return value
        except json.JSONDecodeError:
            pass
    return None


def offsets(namespace, cluster, bootstrap, topic, timeout):
    pod_result = run([
        "kubectl", "-n", namespace, "get", "pod",
        "-l", "strimzi.io/cluster=" + cluster + ",strimzi.io/kind=Kafka",
        "-o", "jsonpath={.items[0].metadata.name}",
    ], timeout)
    pod = pod_result["stdout"].strip()
    if pod_result["exitCode"] != 0 or not pod:
        return {"ok": False, "partitions": [], "errorClass": "broker-pod-unavailable"}
    result = run([
        "kubectl", "-n", namespace, "exec", pod, "--",
        "bin/kafka-get-offsets.sh", "--bootstrap-server", bootstrap,
        "--topic", topic, "--time", "-2",
    ], timeout)
    end_result = run([
        "kubectl", "-n", namespace, "exec", pod, "--",
        "bin/kafka-get-offsets.sh", "--bootstrap-server", bootstrap,
        "--topic", topic, "--time", "-1",
    ], timeout)
    earliest = {}
    end = {}
    for output, target in ((result["stdout"], earliest), (end_result["stdout"], end)):
        for line in output.splitlines():
            parts = line.rsplit(":", 2)
            if len(parts) == 3 and parts[1].isdigit() and parts[2].isdigit():
                target[int(parts[1])] = int(parts[2])
    partitions = [
        {"partition": partition, "earliestOffset": earliest.get(partition), "endOffset": end.get(partition), "retainedRecords": max(0, end.get(partition, 0) - earliest.get(partition, 0))}
        for partition in sorted(set(earliest) | set(end))
    ]
    return {
        "ok": result["exitCode"] == 0 and end_result["exitCode"] == 0,
        "partitions": partitions,
        "errorClass": None if result["exitCode"] == 0 and end_result["exitCode"] == 0 else "offset-query-failed",
    }


def classify(replay, window, execution):
    data = replay.get("data") if isinstance(replay, dict) and isinstance(replay.get("data"), dict) else {}
    source = data.get("source") if isinstance(data.get("source"), dict) else {}
    output = data.get("output") if isinstance(data.get("output"), dict) else {}
    retained = sum(item.get("retainedRecords", 0) for item in window.get("partitions", []))
    scanned = source.get("scanned", 0)
    selected = data.get("selectedFrameCount", 0)
    invalid = source.get("invalidJson", 0)
    ignored = data.get("ignoredFrames") if isinstance(data.get("ignoredFrames"), dict) else {}
    if execution.get("timedOut") or window.get("errorClass") or execution.get("exitCode") in (126, 127):
        return "network-or-runtime-error"
    if isinstance(replay, dict) and replay.get("ok") is True and output.get("recordCount", 0) > 0:
        return "matched"
    if retained == 0 and sum(item.get("endOffset", 0) for item in window.get("partitions", [])) == 0:
        return "producer-not-written"
    if retained == 0:
        return "retention-offset-empty-window"
    error = replay.get("error") if isinstance(replay.get("error"), dict) else {}
    details = error.get("details") if isinstance(error.get("details"), dict) else {}
    if invalid and not selected:
        return "decode-schema-rejection"
    if scanned or ignored or details.get("inputRecordCount") == 0:
        return "filter-mismatch"
    return "replay-rejected"


namespace = os.environ["KAFKA_NAMESPACE"]
cluster = os.environ["KAFKA_CLUSTER"]
bootstrap = os.environ["KAFKA_BOOTSTRAP"]
topic = os.environ["KAFKA_INPUT_TOPIC"]
manager_namespace = os.environ["AGENTRUN_NAMESPACE"]
manager_deployment = os.environ["AGENTRUN_MANAGER_DEPLOYMENT"]
timeout_ms = int(os.environ["KAFKA_REPLAY_TIMEOUT_MS"])
command_grace_ms = int(os.environ["KAFKA_REPLAY_COMMAND_GRACE_MS"])
timeout_seconds = max(1, int((timeout_ms + command_grace_ms) / 1000))
window = offsets(namespace, cluster, bootstrap, topic, timeout_seconds)
command = [
    "kubectl", "-n", manager_namespace, "exec", "deployment/" + manager_deployment, "--",
    "sh", "-c",
    "cd \"${AGENTRUN_APP_ROOT:-/workspace/agentrun}\" && exec ./scripts/agentrun \"$@\"",
    "agentrun", "kafka", "regenerate", "agentrun",
    "--session-id", os.environ["KAFKA_REPLAY_SESSION_ID"],
    "--input-topic", topic,
    "--limit", os.environ["KAFKA_REPLAY_LIMIT"],
    "--timeout-ms", str(timeout_ms),
    "--no-publish", "--format", "json",
]
trace_id = os.environ.get("KAFKA_REPLAY_TRACE_ID", "")
if trace_id:
    command.extend(["--trace-id", trace_id])
execution = run(command, timeout_seconds)
replay = json_object(execution["stdout"]) or json_object(execution["stderr"]) or {}
reason = classify(replay, window, execution)
data = replay.get("data") if isinstance(replay.get("data"), dict) else {}
source = data.get("source") if isinstance(data.get("source"), dict) else {}
output = data.get("output") if isinstance(data.get("output"), dict) else {}
payload = {
    "ok": reason == "matched",
    "networkPlane": "k3s-application-pod",
    "application": "agentrun",
    "namespace": manager_namespace,
    "workload": "deployment/" + manager_deployment,
    "topic": topic,
    "partitionOffsetWindow": window,
    "scanned": source.get("scanned", 0),
    "accepted": data.get("selectedFrameCount", 0),
    "rejected": {
        "invalidJson": source.get("invalidJson", 0),
        "byReason": data.get("ignoredFrames", {}),
        "classification": reason,
    },
    "produced": output.get("recordCount", 0),
    "firstMismatchIdentity": None if reason == "matched" else {"sessionId": os.environ["KAFKA_REPLAY_SESSION_ID"], "traceId": trace_id or None, "reason": reason},
    "mutation": False,
    "topicAppended": False,
    "valuesPrinted": False,
    "execution": {
        "exitCode": execution["exitCode"],
        "elapsedMs": execution["elapsedMs"],
        "timedOut": execution["timedOut"],
        "stderrTail": execution["stderr"][-800:],
    },
    "replay": data,
}
print(json.dumps(payload, ensure_ascii=False))
sys.exit(0)