#!/usr/bin/env python3
"""spark_memory_tuning_validator.py

A small, vendor-neutral helper for validating Apache Spark memory tuning changes.

What it does:
- Reads one or two Spark metrics/log files.
- Extracts common memory-related indicators such as spill counts, GC time, failures,
  and optional executor/driver memory fields when present in the input.
- Produces a concise comparison report that helps decide whether a tuning change
  reduced memory pressure without shifting the bottleneck elsewhere.

What it does not do:
- It does not modify Spark settings.
- It does not connect to a cluster.
- It does not assume any vendor-specific monitoring system.

Recommended use:
1. Export or save relevant Spark logs/metrics after a baseline run.
2. Run the script against the baseline file and the tuned file.
3. Compare spill, GC, and failure indicators before promoting changes.

Supported input:
- Plain text Spark logs
- CSV-like metrics exports
- JSON lines or JSON documents containing metrics fields

Example:
    python spark_memory_tuning_validator.py \
        --before baseline.log \
        --after tuned.log \
        --output report.txt

You can also run with a single file to get a standalone summary:
    python spark_memory_tuning_validator.py --input spark.log
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from dataclasses import dataclass, asdict
from typing import Dict, Iterable, List, Optional, Tuple


METRIC_PATTERNS = {
    "gc_time_ms": [
        re.compile(r"(?i)gc\s*time\s*[:=]\s*(\d+(?:\.\d+)?)\s*ms"),
        re.compile(r"(?i)garbage\s*collection\s*time\s*[:=]\s*(\d+(?:\.\d+)?)\s*ms"),
    ],
    "spill_bytes": [
        re.compile(r"(?i)spill(?:ed)?\s*[:=]\s*(\d+(?:\.\d+)?)\s*(bytes?|kb|mb|gb)?"),
        re.compile(r"(?i)memory\s*spill\s*[:=]\s*(\d+(?:\.\d+)?)\s*(bytes?|kb|mb|gb)?"),
    ],
    "oom_count": [
        re.compile(r"(?i)outofmemoryerror|oom|java\.lang\.outofmemoryerror"),
    ],
    "executor_failed_count": [
        re.compile(r"(?i)executor\s+lost|executor\s+failed|container\s+failed"),
    ],
    "driver_failed_count": [
        re.compile(r"(?i)driver\s+lost|driver\s+failed|driver\s+oom|driver\s+outofmemoryerror"),
    ],
    "shuffle_spill_count": [
        re.compile(r"(?i)shuffle\s*spill"),
        re.compile(r"(?i)spilled\s*to\s*disk"),
    ],
}

UNIT_MULTIPLIERS = {
    None: 1,
    "bytes": 1,
    "byte": 1,
    "kb": 1024,
    "mb": 1024 ** 2,
    "gb": 1024 ** 3,
}


@dataclass
class MetricsSummary:
    source: str
    lines: int = 0
    gc_time_ms: float = 0.0
    spill_bytes: float = 0.0
    oom_count: int = 0
    executor_failed_count: int = 0
    driver_failed_count: int = 0
    shuffle_spill_count: int = 0

    def score(self) -> float:
        """A simple heuristic score: lower is better."""
        return (
            self.gc_time_ms
            + (self.spill_bytes / (1024 ** 2))
            + (self.oom_count * 1000)
            + (self.executor_failed_count * 500)
            + (self.driver_failed_count * 500)
            + (self.shuffle_spill_count * 10)
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Validate Apache Spark memory tuning changes from logs or exported metrics.",
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--input", help="Single log/metrics file to summarize.")
    group.add_argument("--before", help="Baseline log/metrics file.")
    parser.add_argument("--after", help="Tuned log/metrics file to compare against --before.")
    parser.add_argument(
        "--output",
        help="Optional output file for the report. If omitted, prints to stdout.",
    )
    parser.add_argument(
        "--format",
        choices=["text", "json"],
        default="text",
        help="Output format for the report.",
    )
    return parser.parse_args()


def validate_path(path: str) -> None:
    if not path:
        raise ValueError("Empty path provided.")
    if not os.path.isfile(path):
        raise FileNotFoundError(f"File not found: {path}")


def read_lines(path: str) -> List[str]:
    validate_path(path)
    with open(path, "r", encoding="utf-8", errors="replace") as handle:
        return handle.readlines()


def parse_quantity(value: str, unit: Optional[str]) -> float:
    number = float(value)
    return number * UNIT_MULTIPLIERS.get((unit or None).lower() if unit else None, 1)


def extract_metrics(lines: Iterable[str], source: str) -> MetricsSummary:
    summary = MetricsSummary(source=source)
    for line in lines:
        summary.lines += 1
        for pattern in METRIC_PATTERNS["gc_time_ms"]:
            match = pattern.search(line)
            if match:
                summary.gc_time_ms += float(match.group(1))
                break
        for pattern in METRIC_PATTERNS["spill_bytes"]:
            match = pattern.search(line)
            if match:
                summary.spill_bytes += parse_quantity(match.group(1), match.group(2))
                break
        if any(p.search(line) for p in METRIC_PATTERNS["oom_count"]):
            summary.oom_count += 1
        if any(p.search(line) for p in METRIC_PATTERNS["executor_failed_count"]):
            summary.executor_failed_count += 1
        if any(p.search(line) for p in METRIC_PATTERNS["driver_failed_count"]):
            summary.driver_failed_count += 1
        if any(p.search(line) for p in METRIC_PATTERNS["shuffle_spill_count"]):
            summary.shuffle_spill_count += 1
    return summary


def summarize_file(path: str) -> MetricsSummary:
    return extract_metrics(read_lines(path), path)


def compare(before: MetricsSummary, after: MetricsSummary) -> Dict[str, object]:
    def delta(a: float, b: float) -> float:
        return b - a

    return {
        "before": asdict(before),
        "after": asdict(after),
        "delta": {
            "gc_time_ms": delta(before.gc_time_ms, after.gc_time_ms),
            "spill_bytes": delta(before.spill_bytes, after.spill_bytes),
            "oom_count": after.oom_count - before.oom_count,
            "executor_failed_count": after.executor_failed_count - before.executor_failed_count,
            "driver_failed_count": after.driver_failed_count - before.driver_failed_count,
            "shuffle_spill_count": after.shuffle_spill_count - before.shuffle_spill_count,
            "score": after.score() - before.score(),
        },
        "interpretation": interpret_change(before, after),
    }


def interpret_change(before: MetricsSummary, after: MetricsSummary) -> str:
    improved = []
    worsened = []

    if after.oom_count < before.oom_count:
        improved.append("fewer OOM indicators")
    elif after.oom_count > before.oom_count:
        worsened.append("more OOM indicators")

    if after.gc_time_ms < before.gc_time_ms:
        improved.append("less GC time")
    elif after.gc_time_ms > before.gc_time_ms:
        worsened.append("more GC time")

    if after.spill_bytes < before.spill_bytes:
        improved.append("less spill")
    elif after.spill_bytes > before.spill_bytes:
        worsened.append("more spill")

    if after.executor_failed_count < before.executor_failed_count:
        improved.append("fewer executor failures")
    elif after.executor_failed_count > before.executor_failed_count:
        worsened.append("more executor failures")

    if after.driver_failed_count < before.driver_failed_count:
        improved.append("fewer driver failures")
    elif after.driver_failed_count > before.driver_failed_count:
        worsened.append("more driver failures")

    if after.shuffle_spill_count < before.shuffle_spill_count:
        improved.append("fewer shuffle spill events")
    elif after.shuffle_spill_count > before.shuffle_spill_count:
        worsened.append("more shuffle spill events")

    if improved and not worsened:
        return "Likely improvement: " + ", ".join(improved) + "."
    if worsened and not improved:
        return "Likely regression: " + ", ".join(worsened) + "."
    if improved or worsened:
        return (
            "Mixed result: "
            + ("improved in " + ", ".join(improved) if improved else "no clear improvements")
            + "; "
            + ("worsened in " + ", ".join(worsened) if worsened else "no clear regressions")
            + "."
        )
    return "No clear memory-related signal detected in the provided input."


def format_text_report(payload: Dict[str, object]) -> str:
    if "after" not in payload:
        summary = payload["summary"]
        return (
            f"Source: {summary['source']}\n"
            f"Lines: {summary['lines']}\n"
            f"GC time (ms): {summary['gc_time_ms']}\n"
            f"Spill (bytes): {summary['spill_bytes']}\n"
            f"OOM indicators: {summary['oom_count']}\n"
            f"Executor failures: {summary['executor_failed_count']}\n"
            f"Driver failures: {summary['driver_failed_count']}\n"
            f"Shuffle spill events: {summary['shuffle_spill_count']}\n"
            f"Score: {summary['score']:.2f}\n"
        )

    before = payload["before"]
    after = payload["after"]
    delta = payload["delta"]
    return (
        "Spark Memory Tuning Validation Report\n"
        "-------------------------------------\n"
        f"Before: {before['source']}\n"
        f"After:  {after['source']}\n\n"
        f"GC time (ms): {before['gc_time_ms']} -> {after['gc_time_ms']} (delta {delta['gc_time_ms']:+.2f})\n"
        f"Spill (bytes): {before['spill_bytes']} -> {after['spill_bytes']} (delta {delta['spill_bytes']:+.2f})\n"
        f"OOM indicators: {before['oom_count']} -> {after['oom_count']} (delta {delta['oom_count']:+d})\n"
        f"Executor failures: {before['executor_failed_count']} -> {after['executor_failed_count']} (delta {delta['executor_failed_count']:+d})\n"
        f"Driver failures: {before['driver_failed_count']} -> {after['driver_failed_count']} (delta {delta['driver_failed_count']:+d})\n"
        f"Shuffle spill events: {before['shuffle_spill_count']} -> {after['shuffle_spill_count']} (delta {delta['shuffle_spill_count']:+d})\n"
        f"Score: {before['score']:.2f} -> {after['score']:.2f} (delta {delta['score']:+.2f})\n\n"
        f"Interpretation: {payload['interpretation']}\n"
    )


def main() -> int:
    args = parse_args()
    try:
        if args.input:
            summary = summarize_file(args.input)
            payload = {"summary": asdict(summary)}
        else:
            if not args.after:
                raise ValueError("--after is required when using --before.")
            before = summarize_file(args.before)
            after = summarize_file(args.after)
            payload = compare(before, after)

        output = json.dumps(payload, indent=2) if args.format == "json" else format_text_report(payload)

        if args.output:
            with open(args.output, "w", encoding="utf-8") as handle:
                handle.write(output)
                if not output.endswith("\n"):
                    handle.write("\n")
        else:
            print(output)
        return 0
    except Exception as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())