#!/usr/bin/env python3
import argparse
import io
import json
import re
from collections import defaultdict
from pathlib import Path

import pandas as pd
from rich.console import Console
from rich.markdown import Markdown
from XRootD.client import FileSystem, File

def save_markdown_to_file(output, filename="output.md"):
    """Saves a list of strings as Markdown to a file."""
    try:
        with open(filename, "w", encoding="utf-8") as f:
            f.write("\n".join(output))
        print(f"Markdown output saved to {filename}")
    except Exception as e:
        print(f"Error saving Markdown to file: {e}")

def to_3sf(x):
    """Convert a number to a string with 3 significant figures."""
    if isinstance(x, str):
        return x
    return f"{float(f'{x:.3g}'):.6g}"


instance = "root://eoslhcb.cern.ch"
base_dir = "/eos/lhcb/storage/lhcbpr/www/UpgradeRateTest"
spruce_stream_definitions = {
    "b2oc": ["SpruceB2OC.*", "SpruceCalib.*"],
    "bandq": [ "SpruceBandQ.*"],
    "b2cc": ["SpruceB2CC.*"],
    "qee": ["SpruceQEE.*"],
    "rd": ["SpruceRD.*"],
    "bnoc": ["SpruceBnoC.*"],
    "ift": ["SpruceIFT.*"],
    "sl": ["SpruceSL.*"],

}
columns = ["Line", "Total Retention (%)", "Rate (kHz)", "Avg DstData Size (kB)", "Total Bandwidth (GB/s)"]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", type=Path, required=False)
    args = parser.parse_args()

    if args.csv:
        latest_name = None
        raw_csv = args.csv.read_bytes()
    else:
        eos = FileSystem(instance)
        status, listing = eos.dirlist(base_dir)
        if not status.ok:
            raise NotImplementedError(status)
        latest_idx = None
        latest_name = None
        for l in listing:
            if match := re.search(r"BandwidthTest_lhcb-master.(\d+)_Moore_hlt2_and_spruce_bandwidth_x86_64_v3-el9-gcc\d+-opt\+g_(.+)", l.name):
                if latest_idx and latest_idx > int(match.group(1)):
                    continue
                latest_idx = int(match.group(1))
                latest_name = l.name
        f = File()
        status, _ = f.open(f"{instance}/{base_dir}/{latest_name}/spruce__full__rates_for_all_lines.csv")
        if not status.ok:
            raise NotImplementedError(status)
        status, raw_csv = f.read()
        if not status.ok:
            raise NotImplementedError(status)

    df = pd.read_csv(io.BytesIO(raw_csv))

    stream_to_lines = defaultdict(list)
    threshold = defaultdict(lambda: 0.001)
    threshold["trackeff"] = -1
    really_low_threshold = -1
    for stream, patterns in spruce_stream_definitions.items():
        stream_to_lines[stream] = []
        stream_to_lines[f"{stream}low"] = []
        for _, r in df.iterrows():
            if r["Total Bandwidth (GB/s)"] > threshold[stream]:
                lines = stream_to_lines[stream]
            elif r["Total Bandwidth (GB/s)"] > really_low_threshold:
                lines = stream_to_lines[f"{stream}low"]
            else:
                lines = stream_to_lines["reallylow"]
            if any(re.fullmatch(p, r["Line"]) for p in patterns):
                lines.append(r["Line"])
    stream_to_lines = {k: v for k, v in stream_to_lines.items() if v}

    ## Check no lines duplicated btw streams
    assert sum(map(len, stream_to_lines.values())) == len({x for y in stream_to_lines.values() for x in y})
    output = ["# Stream split proposal\n"]
    if latest_name:
        output.append(f"Based on: https://lhcbpr-hlt.web.cern.ch/UpgradeRateTest/{latest_name}/hlt2__production__all_rates.html")
    else:
        output.append("Based on local test data")
    output += [
        "Parameters:\n",
        f" * {dict(threshold)=}",
        f" * {really_low_threshold=}",
    ]
    for name, lines in stream_to_lines.items():
        output += [
            f"\n## `{name.upper()}` ({len(lines)}) lines\n",
            "<details><summary>Click to expand</summary>\n",
            "| " + " | ".join(columns) + " |",
            "| " + " | ".join(["---"]*len(columns)) + " |"
        ]
        output += [
            "| " + " | ".join(map(to_3sf, df.loc[df["Line"] == line, columns].iloc[0])) + " |"
            for line in lines
        ]
        output.append("</details>\n")

    output.append("## Crude bandwidth estimates (DO NOT TRUST! Double counts overlaps!)\n")
    by_stream = {}
    for name, lines in stream_to_lines.items():
        by_stream[name] = sum(df.loc[df["Line"] == line, columns].iloc[0]["Total Bandwidth (GB/s)"] for line in lines)
    total = sum(by_stream.values())
    for name, value in by_stream.items():
        if name.endswith("low"):
            continue
        line = f"* `{name}`: {to_3sf(value)} GB/s ({value/total:.1%})"
        if f"{name}low" in by_stream:
            line += f" + {to_3sf(by_stream[f'{name}low'])} GB/s ({by_stream[f'{name}low']/total:.1%})"
            line += f" (high/low ratio {by_stream[f'{name}low']/value:.1%})"
        output.append(line)

    Path("stream_to_lines.json").write_text(json.dumps(stream_to_lines, indent=4))
    Path("stream_to_lines.md").write_text("\n".join(output))

    save_markdown_to_file(output, "FULL_restreaming.md")
    Console().print(Markdown("\n".join(output)))


if __name__ == "__main__":
    main()