import subprocess
import json
import sys
import os
import time
import tempfile
import base64
from datetime import datetime
from pathlib import Path
from warehouse import Media_flag, rand12

Recordings = Path.home() / "records"
Tmp = tempfile.gettempdir()

def run(cmd):
    return subprocess.Popen(c, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

def find_focused_pid(node):
    if isinstance(node, dict):
        if node.get("focused") is True:
            return node.get("pid")

        for key in ("nodes", "floating_nodes"):
            for child in node.get(key, []):
                pid = find_focused_pid(child)
                if pid:
                    return pid

    elif isinstance(node, list):
        for item in node:
            pid = find_focused_pid(item)
            if pid:
                return pid

    return None

def force_kill():
    tree = json.loads(subprocess.check_output(
        ["swaymsg", "-t", "get_tree"],
        text=True
    ))
    pid = find_focused_pid(tree)

    if not pid:
        return

    os.kill(pid, 9)

def is_recording():
    return subprocess.run(
        ["pgrep", "-x", "wf-recorder"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL
    ).returncode == 0


def record(fullscreen=False):
    Recordings.mkdir(parents=True, exist_ok=True)

    if is_recording():
        subprocess.run(["killall", "wf-recorder"])
        try:
            os.remove(Media_flag)
        except FileNotFoundError:
            pass
        return

    Path(Media_flag).write_text("1")

    slurp_cmd = ["slurp"]
    if fullscreen:
        slurp_cmd.append("-o")

    selection = subprocess.getoutput(" ".join(slurp_cmd))

    if not selection.strip():
        return

    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    output = Recordings / f"{timestamp}.mp4"

    proc = subprocess.Popen(
        ["wf-recorder", "-g", selection, "-f", str(output)]
    )
    proc.wait()

    if not output.exists():
        return

    compressed = output.with_name(output.stem + "_comp720p.mp4")

    subprocess.run([
        "ffmpeg",
        "-i", str(output),
        "-vf", "scale=-2:720",
        "-vcodec", "libx264",
        "-crf", "32",
        "-preset", "fast",
        str(compressed)
    ])

    try:
        output.unlink()
    except FileNotFoundError:
        pass

def resolve_output(out):
    if not out:
        out = Tmp

    out = os.path.expanduser(out)

    if os.path.isdir(out) or out.endswith("/"):
        os.makedirs(out, exist_ok=True)
        return os.path.join(out, f"{rand12()}.png")

    parent = os.path.dirname(out)
    if parent:
        os.makedirs(parent, exist_ok=True)

    if not out.endswith(".png"):
        out += ".png"

    return out


def take_screenshot(output_path):
    Path(Media_flag).write_text("1")

    try:
        time.sleep(0.2)

        freeze = subprocess.Popen(["wayfreeze"])
        time.sleep(0.1)

        try:
            region = subprocess.check_output(["slurp"]).decode().strip()

            subprocess.run(
                ["grim", "-g", region, output_path],
                check=True
            )

        except subprocess.CalledProcessError:
            return False

        finally:
            freeze.terminate()

        if os.path.exists(output_path):
            with open(output_path, "rb") as f:
                subprocess.run(
                    ["wl-copy", "--type", "image/png"],
                    input=f.read()
                )

            subprocess.run(["notify-send", "screenshot", f"Saved: {output_path}"])
            return True

    finally:
        try:
            os.remove(Media_flag)
        except FileNotFoundError:
            pass

    return False

def main():
    if len(sys.argv) < 2:
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "kill":
        force_kill()

    elif cmd == "record":
        fullscreen = "--fullscreen" in sys.argv
        record(fullscreen)

    elif cmd == "screenshot":
        out = None

        if "--out" in sys.argv:
            idx = sys.argv.index("--out")
            if idx + 1 < len(sys.argv):
                out = sys.argv[idx + 1]

        output = resolve_output(out)
        take_screenshot(output)

if __name__ == "__main__":
    main()