#!/usr/bin/env python3

import os
import re
import sys
import urllib.parse
from pathlib import Path

import requests
from cryptography.hazmat.primitives.ciphers.aead import AESGCM


home = Path.home()
if (home / "downloads").exists():
    DOWNLOAD_DIR = home / "downloads" / "xmpp"
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
else:
    DOWNLOAD_DIR = home / "Downloads" / "xmpp"
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)

def log(msg: str) -> None:
    log_file = Path.home() / ".local" / "share" / "profanity-aesgcm" / "handler.log"
    log_file.parent.mkdir(parents=True, exist_ok=True)
    with log_file.open("a", encoding="utf-8") as f:
        f.write(msg + "\n")


def extract_uri(message: str) -> str | None:
    m = re.search(r'(aesgcm://\S+)', message)
    return m.group(1) if m else None


def parse_aesgcm_uri(uri: str):
    parsed = urllib.parse.urlparse(uri)

    if parsed.scheme != "aesgcm":
        raise ValueError("Unsupported scheme")

    fragment = parsed.fragment.strip()
    frag_hex = "".join(fragment.split())  # remove accidental spaces/newlines

    # XEP-0454 style: 12-byte IV (24 hex chars) + 32-byte key (64 hex chars)
    if len(frag_hex) < 88:
        raise ValueError(f"Fragment too short: {len(frag_hex)} hex chars")

    iv_hex = frag_hex[:24]
    key_hex = frag_hex[24:88]

    iv = bytes.fromhex(iv_hex)
    key = bytes.fromhex(key_hex)

    https_url = urllib.parse.urlunparse(("https", parsed.netloc, parsed.path, "", "", ""))
    filename = os.path.basename(parsed.path) or "download.bin"

    return https_url, filename, key, iv


def download_file(url: str) -> bytes:
    r = requests.get(url, timeout=60)
    r.raise_for_status()
    return r.content


def decrypt_xmpp_aesgcm(data: bytes, key: bytes, iv: bytes) -> bytes:
    # AESGCM.decrypt expects ciphertext || tag
    aesgcm = AESGCM(key)
    return aesgcm.decrypt(iv, data, None)


def unique_path(path: Path) -> Path:
    if not path.exists():
        return path

    stem = path.stem
    suffix = path.suffix

    for i in range(1, 1000):
        candidate = path.with_name(f"{stem}-{i}{suffix}")
        if not candidate.exists():
            return candidate

    raise RuntimeError("Could not create unique filename")

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

    message = sys.argv[1]
    uri = extract_uri(message)
    if not uri:
        sys.exit(0)

    try:
        https_url, filename, key, iv = parse_aesgcm_uri(uri)
        enc = download_file(https_url)
        plain = decrypt_xmpp_aesgcm(enc, key, iv)

        out_path = unique_path(DOWNLOAD_DIR / filename)
        out_path.write_bytes(plain)

        log(f"OK saved: {out_path}")
        print(str(out_path))

        # Optional: uncomment if you want auto-open on desktop
        # os.system(f'xdg-open "{out_path}" >/dev/null 2>&1 &')

    except Exception as e:
        log(f"ERROR: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()