import os
import re
import sys

import matplotlib.pyplot as plt

OUT_DIR = "out"
os.makedirs(OUT_DIR, exist_ok=True)

LOG_RE = re.compile(
    r"(PROC_START|PROC_END)\s+pid=(\d+)\s+ppid=(\d+)\s+depth=(\d+)\s+l=(\d+)\s+r=(\d+)\s+ts=([\d.]+)"
)


class Proc:
    def __init__(self, pid, ppid, depth):
        self.pid = pid
        self.ppid = ppid
        self.depth = depth
        self.start = None
        self.end = None
        self.children = []

    @property
    def duration(self):
        return self.end - self.start


def parse_log(path):
    procs = {}

    for line in open(path):
        m = LOG_RE.search(line)
        if not m:
            continue

        typ, pid, ppid, depth, l, r, ts = m.groups()
        pid, ppid, depth, ts = int(pid), int(ppid), int(depth), float(ts)

        if pid not in procs:
            procs[pid] = Proc(pid, ppid, depth)

        if typ == "PROC_START":
            procs[pid].start = ts
        else:
            procs[pid].end = ts

    return procs


def build_tree(procs):
    root = None
    for p in procs.values():
        if p.ppid in procs:
            procs[p.ppid].children.append(p)
        else:
            root = p
    return root


def dfs_order(root):
    order = []

    def dfs(p):
        order.append(p)
        for c in sorted(p.children, key=lambda x: x.start):
            dfs(c)

    dfs(root)
    return order


def write_md(procs, ordered):
    out_path = os.path.join(OUT_DIR, "processes.md")
    first = min(p.start for p in procs.values())

    with open(out_path, "w") as f:
        f.write("| PID | PPID | Depth | Start | End | Offset | Duration |\n")
        f.write("|---|---|---|---|---|---|---|\n")

        for p in ordered:
            f.write(
                f"| {p.pid} | {p.ppid} | {p.depth} | "
                f"{p.start:.6f} | {p.end:.6f} | "
                f"{p.start - first:.6f} | {p.duration:.6f} |\n"
            )


def draw_gantt(ordered):
    out_tree = os.path.join(OUT_DIR, "gantt_tree.png")
    out_pid = os.path.join(OUT_DIR, "gantt_pid.png")

    first = min(p.start for p in ordered)

    # -------------------------
    # 1) TREE ORDER (current)
    # -------------------------
    fig, ax = plt.subplots(figsize=(12, 6))

    for i, p in enumerate(ordered):
        ax.barh(i, p.duration, left=p.start - first, height=0.6)
        ax.text(p.start - first, i, f"PID {p.pid} ({p.depth})", va="center")

    ax.set_yticks(range(len(ordered)))
    ax.set_yticklabels([f"{p.pid} ({p.depth})" for p in ordered])
    ax.set_xlabel("Seconds from first process start")
    ax.set_title("Gantt (Tree Order)")
    ax.invert_yaxis()

    plt.tight_layout()
    plt.savefig(out_tree)
    plt.close()

    # -------------------------
    # 2) PID-SORTED VIEW
    # -------------------------
    pid_order = sorted(ordered, key=lambda p: p.pid)

    fig, ax = plt.subplots(figsize=(12, 6))

    for i, p in enumerate(pid_order):
        ax.barh(i, p.duration, left=p.start - first, height=0.6)
        ax.text(p.start - first, i, f"PID {p.pid} ({p.depth})", va="center")

    ax.set_yticks(range(len(pid_order)))
    ax.set_yticklabels([f"{p.pid} ({p.depth})" for p in pid_order])
    ax.set_xlabel("Seconds from first process start")
    ax.set_title("Gantt (PID Order)")
    ax.invert_yaxis()

    plt.tight_layout()
    plt.savefig(out_pid)
    plt.close()


def main():
    procs = parse_log(sys.argv[1])
    root = build_tree(procs)
    ordered = dfs_order(root)

    write_md(procs, ordered)
    draw_gantt(ordered)

    print(f"Outputs saved to: {OUT_DIR}/")


if __name__ == "__main__":
    main()