#!/usr/bin/env python3
import sys
import base64
import zipfile
import io
import json
import re
import requests

BUILD_NUMBER = sys.argv[1]
HTML_FILE    = sys.argv[2]
PUSHGATEWAY  = "http://100.106.10.63:9091"

# Step 1: Extract base64 ZIP from the HTML
with open(HTML_FILE, "r", encoding="utf-8") as f:
    content = f.read()

match = re.search(
    r'<script id="playwrightReportBase64"[^>]*>(data:application/zip;base64,([^<]+))</script>',
    content
)
if not match:
    print("ERROR: Could not find embedded report data in HTML")
    sys.exit(1)

b64_data = match.group(2).strip()
zip_bytes = base64.b64decode(b64_data)

# Step 2: Open ZIP and read all JSON files
with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
    with zf.open("report.json") as jf:
        report = json.load(jf)

    test_files = {}
    for name in zf.namelist():
        if name != "report.json" and name.endswith(".json"):
            with zf.open(name) as jf:
                test_files[name] = json.load(jf)

# Step 3: Extract top-level stats
stats       = report.get("stats", {})
total       = stats.get("total",      0)
passed      = stats.get("expected",   0)
failed      = stats.get("unexpected", 0)
flaky       = stats.get("flaky",      0)
skipped     = stats.get("skipped",    0)
duration_ms = stats.get("duration",   0)
start_time  = stats.get("startTime",  0)

print(f"\n=== Build {BUILD_NUMBER} Summary ===")
print(f"  Total:    {total}")
print(f"  Passed:   {passed}")
print(f"  Failed:   {failed}")
print(f"  Flaky:    {flaky}")
print(f"  Skipped:  {skipped}")
print(f"  Duration: {duration_ms}ms")

# Step 4: Build metrics payload correctly
# Group by metric name - HELP and TYPE only once per metric
metrics = {}

def add_metric(name, help_text, metric_type, value, labels):
    if name not in metrics:
        metrics[name] = {
            "help": help_text,
            "type": metric_type,
            "samples": []
        }
    label_str = ",".join([f'{k}="{v}"' for k, v in labels.items()])
    metrics[name]["samples"].append(f"{name}{{{label_str}}} {value}")

# --- Top-level summary metrics ---
add_metric("playwright_tests_total",   "Total number of tests",           "gauge", total,       {"build": BUILD_NUMBER})
add_metric("playwright_tests_passed",  "Passed tests",                    "gauge", passed,      {"build": BUILD_NUMBER})
add_metric("playwright_tests_failed",  "Failed tests",                    "gauge", failed,      {"build": BUILD_NUMBER})
add_metric("playwright_tests_flaky",   "Flaky tests",                     "gauge", flaky,       {"build": BUILD_NUMBER})
add_metric("playwright_tests_skipped", "Skipped tests",                   "gauge", skipped,     {"build": BUILD_NUMBER})
add_metric("playwright_duration_ms",   "Total test suite duration in ms", "gauge", duration_ms, {"build": BUILD_NUMBER})
add_metric("playwright_start_time",    "Test suite start time epoch ms",  "gauge", start_time,  {"build": BUILD_NUMBER})

# --- Pass rate ---
pass_rate = round((passed / total * 100), 2) if total > 0 else 0
add_metric("playwright_pass_rate", "Pass rate percentage", "gauge", pass_rate, {"build": BUILD_NUMBER})

# --- Per-file/suite metrics ---
for file_entry in report.get("files", []):
    file_name  = file_entry.get("fileName", "unknown").replace('"', '').replace(' ', '_')
    file_stats = file_entry.get("stats", {})

    add_metric("playwright_suite_total",       "Total tests per file",    "gauge", file_stats.get("total", 0),      {"build": BUILD_NUMBER, "file": file_name})
    add_metric("playwright_suite_passed",      "Passed tests per file",   "gauge", file_stats.get("expected", 0),   {"build": BUILD_NUMBER, "file": file_name})
    add_metric("playwright_suite_failed",      "Failed tests per file",   "gauge", file_stats.get("unexpected", 0), {"build": BUILD_NUMBER, "file": file_name})
    add_metric("playwright_suite_flaky",       "Flaky tests per file",    "gauge", file_stats.get("flaky", 0),      {"build": BUILD_NUMBER, "file": file_name})
    add_metric("playwright_suite_skipped",     "Skipped tests per file",  "gauge", file_stats.get("skipped", 0),    {"build": BUILD_NUMBER, "file": file_name})
    add_metric("playwright_suite_duration_ms", "Duration per file in ms", "gauge", file_stats.get("duration", 0),   {"build": BUILD_NUMBER, "file": file_name})

# --- Per-test metrics ---
for file_name, file_data in test_files.items():
    for test in file_data.get("tests", []):
        test_title  = test.get("title", "unknown").replace('"', '').replace('\n', ' ').replace('\\', '')[:80]
        test_file   = test.get("location", {}).get("file", "unknown").replace('"', '')
        test_line   = str(test.get("location", {}).get("line", 0))
        project     = test.get("projectName", "default").replace('"', '')
        outcome     = test.get("outcome", "unknown")
        duration    = test.get("duration", 0)
        retry_count = max(len(test.get("results", [])) - 1, 0)

        outcome_map = {"expected": 1, "unexpected": 0, "flaky": 2, "skipped": 3}
        outcome_num = outcome_map.get(outcome, -1)

        labels = {
            "build":   BUILD_NUMBER,
            "test":    test_title,
            "file":    test_file,
            "line":    test_line,
            "project": project,
        }

        add_metric("playwright_test_duration_ms", "Individual test duration ms",               "gauge", duration,    labels)
        add_metric("playwright_test_outcome",     "Test outcome 1=pass 0=fail 2=flaky 3=skip", "gauge", outcome_num, labels)
        add_metric("playwright_test_retries",     "Number of retries",                         "gauge", retry_count, labels)

        print(f"  [{outcome.upper():10}] {test_title[:60]} ({duration}ms, retries={retry_count})")

# Step 5: Build final payload - HELP and TYPE only once per metric
payload_lines = []
for name, data in metrics.items():
    payload_lines.append(f"# HELP {name} {data['help']}")
    payload_lines.append(f"# TYPE {name} {data['type']}")
    payload_lines.extend(data["samples"])

metrics_payload = "\n".join(payload_lines) + "\n"

# Step 6: Push to Pushgateway
response = requests.post(
    f"{PUSHGATEWAY}/metrics/job/playwright/instance/build-{BUILD_NUMBER}",
    data=metrics_payload,
    headers={"Content-Type": "text/plain; version=0.0.4"}
)
print(f"\nPushed {len(payload_lines)} lines to Pushgateway → HTTP {response.status_code}")
if response.status_code != 200:
    print(f"Error: {response.text[:200]}")