#!/usr/bin/env python3
import json
import subprocess
from pathlib import Path

out = Path("./mute/monitors.conf")

DIRECT = {
    "left": (-1, 0),
    "right": (1, 0),
    "above": (0, -1),
    "below": (0, 1),
    "primary": (0, 0)
}

def get_sway_data(target_type):
    cmd = ["swaymsg", "-t", f"get_{target_type}", "-r"]
    return json.loads(subprocess.check_output(cmd, text=True))

def get_best_mode(output):
    modes = output.get("modes", [])
    if not modes:
        return (1920, 1080)

    for m in modes:
        if m.get("preferred"):
            return m["width"], m["height"]

    best = max(modes, key=lambda m: m["width"] * m["height"])
    return best["width"], best["height"]

def ask(prompt, default=None):
    prompt_str = f"{prompt} [{default}]: " if default else f"{prompt}: "
    val = input(prompt_str).strip()
    return val if val else default

def compute_positions(monitors, layout, outputs):
    positions = {}
    output_map = {o["name"]: o for o in outputs}

    primary_name = next((name for name, dist in layout.items() if dist == "primary"), None)
    if not primary_name and monitors:
        primary_name = monitors[0]["name"]
        layout[primary_name] = "primary"

    logical_sizes = {}

    for m in monitors:
        name = m["name"]

        try:
            res_w, res_h = map(int, m["res"].lower().split("x"))
        except ValueError:
            res_w, res_h = get_best_mode(output_map[name])

        scale = float(m["scale"])

        logical_sizes[name] = {
            "width": res_w / scale,
            "height": res_h / scale
        }

    for m in monitors:
        name = m["name"]
        direction = layout[name]

        if direction == "primary":
            positions[name] = (0, 0)
            continue

        dx, dy = DIRECT.get(direction, (1, 0))

        ref_w = logical_sizes[name]["width"] if dx < 0 else logical_sizes[primary_name]["width"]
        ref_h = logical_sizes[name]["height"] if dy < 0 else logical_sizes[primary_name]["height"]

        positions[name] = (int(ref_w * dx), int(ref_h * dy))

    return positions

def build_config(monitors, positions):
    lines = ["# generated config\n"]

    for m in monitors:
        name = m["name"]
        mode = f"{m['res']}@{m['refresh']}" if m["refresh"] else m["res"]
        x, y = positions[name]

        lines.append(
            f"output {name} mode {mode}\n"
            f"output {name} scale {m['scale']}\n"
            f"output {name} position {x} {y}\n"
        )

    return "\n".join(lines)

def main():
    Path("mute").mkdir(exist_ok=True)

    outputs = get_sway_data("outputs")
    workspaces = get_sway_data("workspaces")

    monitors = []
    layout_map = {}

    for o in outputs:
        name = o["name"]
        ws = [str(w["name"]) for w in workspaces if w.get("output") == name]

        w, h = get_best_mode(o)
        best_res = f"{w}x{h}"

        print(f"\nMONITOR: {name}\nWORKSPACES: {ws}")

        layout_map[name] = ask(
            "Position relative to primary (primary/left/right/above/below)",
            "right"
        )

        res_input = ask("Resolution (optional press enter for best)", "")

        monitors.append({
            "name": name,
            "res": res_input if res_input else best_res,
            "refresh": ask("Refresh rate (optional)", ""),
            "scale": ask("Scale", "1")
        })

    positions = compute_positions(monitors, layout_map, outputs)
    out.write_text(build_config(monitors, positions))

if __name__ == "__main__":
    main()