#!/usr/bin/env python3
"""Save linked Tana media beside an extracted Markdown export."""

import hashlib
import os
import re
import shutil
import sys
from pathlib import Path
from urllib.parse import unquote, urlsplit
from urllib.request import urlopen


MEDIA_HOSTS = {"firebasestorage.googleapis.com", "storage.googleapis.com"}
MARKDOWN_LINK = re.compile(r"(?P<start>\]\()(?P<url>https://[^)\s]+)(?P<end>\))")


def main() -> int:
    if len(sys.argv) != 2:
        print("Usage: python3 localize-tana-media.py EXTRACTED_EXPORT_FOLDER", file=sys.stderr)
        return 2

    root = Path(sys.argv[1]).expanduser().resolve()
    if not root.is_dir():
        print(f"Not a folder: {root}", file=sys.stderr)
        return 2

    media_dir = root / "_media"
    markdown_files = list(root.rglob("*.md"))
    local_paths = {}
    downloaded = 0
    updated = 0

    for markdown_file in markdown_files:
        original = markdown_file.read_text(encoding="utf-8")

        def replace_link(match):
            nonlocal downloaded
            url = match.group("url")
            parsed = urlsplit(url)
            if parsed.hostname not in MEDIA_HOSTS and not (parsed.hostname or "").endswith(".storage.googleapis.com"):
                return match.group(0)

            if url not in local_paths:
                filename = unquote(parsed.path).rsplit("/", 1)[-1]
                extension = Path(filename).suffix.lower()
                if not re.fullmatch(r"\.[a-z0-9]{1,10}", extension):
                    extension = ".bin"
                digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:20]
                destination = media_dir / f"{digest}{extension}"
                media_dir.mkdir(exist_ok=True)
                temporary = destination.with_name(destination.name + ".part")
                try:
                    with urlopen(url, timeout=60) as response, temporary.open("wb") as output:
                        shutil.copyfileobj(response, output)
                    temporary.replace(destination)
                    local_paths[url] = destination
                    downloaded += 1
                except Exception as error:
                    temporary.unlink(missing_ok=True)
                    local_paths[url] = None
                    print(f"Could not download {parsed.netloc}{parsed.path}: {error}", file=sys.stderr)

            destination = local_paths[url]
            if destination is None:
                return match.group(0)
            relative_path = Path(os.path.relpath(destination, markdown_file.parent)).as_posix()
            return f'{match.group("start")}{relative_path}{match.group("end")}'

        localized = MARKDOWN_LINK.sub(replace_link, original)
        if localized != original:
            markdown_file.write_text(localized, encoding="utf-8")
            updated += 1

    failed = sum(path is None for path in local_paths.values())
    file_label = "file" if updated == 1 else "files"
    print(f"Downloaded {downloaded} media files; updated {updated} Markdown {file_label}; {failed} links failed.")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())