#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["httpx>=0.27,<1"]
# requires-python = ">=3.10"
# ///
"""Upload images to Immich and derive albums and tags from their filenames."""

from __future__ import annotations

import argparse
import hashlib
import logging
import os
import re
import shutil
import struct
import sys
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Any, Sequence

import httpx


LOG = logging.getLogger("immichup")
TAG_PORTRAIT = "aspect:portrait"
TAG_LANDSCAPE = "aspect:landscape"
STUDIO_RE = re.compile(r"\b@(\S+)\b")


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Upload JPEG images to Immich; filename prefixes become albums and tags."
    )
    parser.add_argument(
        "--api-key", default=os.getenv("IMMICH_API_KEY", ""),
        help="Immich API key (defaults to IMMICH_API_KEY)",
    )
    parser.add_argument(
        "--server-url", default=os.getenv("IMMICH_SERVER_URL", ""),
        help="Immich server URL (defaults to IMMICH_SERVER_URL)",
    )
    parser.add_argument(
        "--move-uploaded-to", metavar="DIR",
        help="move successfully uploaded images into this directory",
    )
    parser.add_argument("--debug", action="store_true", help="enable debug logging")
    parser.add_argument(
        "--dry-run", action="store_true", help="show planned uploads without uploading",
    )
    parser.add_argument("images", nargs="+", type=Path, help="image paths to process")
    args = parser.parse_args(argv)
    if not args.dry_run and not args.server_url:
        parser.error("--server-url or IMMICH_SERVER_URL is required unless --dry-run is used")
    if not args.dry_run and not args.api_key:
        parser.error("--api-key or IMMICH_API_KEY is required unless --dry-run is used")
    return args


def parse_filename(image_path: Path) -> tuple[list[str], str | None] | None:
    filename = image_path.stem
    if " -- " not in filename:
        return None
    people = [name.strip() for name in filename.split(" -- ", 1)[0].split(",")]
    if not people:
        return None
    match = STUDIO_RE.search(filename)
    return people, match.group(1) if match else None


def read_jpeg_dimensions(path: Path) -> tuple[int, int]:
    """Read width and height from a JPEG frame header without loading the image."""
    with path.open("rb") as image:
        if image.read(2) != b"\xff\xd8":
            raise ValueError("not a valid JPEG file")
        while True:
            byte = image.read(1)
            if not byte:
                raise ValueError("JPEG has no frame header")
            if byte != b"\xff":
                raise ValueError("invalid JPEG format")
            marker_byte = image.read(1)
            while marker_byte == b"\xff":
                marker_byte = image.read(1)
            if not marker_byte:
                raise ValueError("truncated JPEG marker")
            marker = marker_byte[0]
            if marker in (0xD8, 0xD9) or 0xD0 <= marker <= 0xD7 or marker == 0x01:
                continue
            length_data = image.read(2)
            if len(length_data) != 2:
                raise ValueError("truncated JPEG segment")
            length = struct.unpack(">H", length_data)[0]
            if length < 2:
                raise ValueError("invalid JPEG segment length")
            if 0xC0 <= marker <= 0xC3:
                frame = image.read(length - 2)
                if len(frame) < 5:
                    raise ValueError("truncated JPEG frame header")
                height, width = struct.unpack(">HH", frame[1:5])
                return width, height
            image.seek(length - 2, os.SEEK_CUR)


class ImmichClient:
    def __init__(self, server_url: str, api_key: str) -> None:
        self.http = httpx.Client(
            base_url=server_url.rstrip("/") + "/api/",
            headers={"x-api-key": api_key},
            timeout=60.0,
        )
        self.device_id = str(uuid.uuid4())
        self.album_ids: dict[str, str] = {}
        self.tag_ids: dict[str, str] = {}
        self.metadata_lock = threading.RLock()

    def close(self) -> None:
        self.http.close()

    def _check(self, response: httpx.Response, action: str) -> None:
        if not response.is_success:
            raise RuntimeError(f"{action}: HTTP {response.status_code} {response.reason_phrase}")

    def upload_asset(self, path: Path) -> dict[str, Any]:
        stat = path.stat()
        with path.open("rb") as image:
            digest = hashlib.sha1()
            for chunk in iter(lambda: image.read(1024 * 1024), b""):
                digest.update(chunk)
            checksum = digest.hexdigest()
        modified = datetime.fromtimestamp(stat.st_mtime).astimezone().isoformat()
        form = {
            "deviceAssetId": f"{path.name}-{stat.st_size}",
            "deviceId": self.device_id,
            "assetType": "image",
            "fileCreatedAt": modified,
            "fileModifiedAt": modified,
            "isFavorite": "false",
            "fileExtension": path.suffix.lower(),
            "duration": "0",
            "isReadOnly": "false",
            "isArchived": "false",
        }
        with path.open("rb") as image:
            response = self.http.post(
                "assets",
                headers={"x-immich-checksum": checksum},
                data=form,
                files={"assetData": (path.name, image)},
            )
        self._check(response, "failed to upload asset")
        return response.json()

    def get_or_create_album(self, name: str) -> str:
        with self.metadata_lock:
            if name in self.album_ids:
                return self.album_ids[name]
            response = self.http.get("albums")
            self._check(response, "failed to get albums")
            for album in response.json():
                self.album_ids[album["albumName"]] = album["id"]
            if name not in self.album_ids:
                response = self.http.post("albums", json={"albumName": name})
                self._check(response, "failed to create album")
                album = response.json()
                self.album_ids[name] = album["id"]
            return self.album_ids[name]

    def create_tags(self, names: list[str]) -> dict[str, str]:
        with self.metadata_lock:
            missing = list(dict.fromkeys(name for name in names if name not in self.tag_ids))
            if missing:
                response = self.http.put("tags", json={"tags": missing})
                self._check(response, "failed to create tags")
                self.tag_ids.update({tag["name"]: tag["id"] for tag in response.json()})
            return {name: self.tag_ids[name] for name in names if name in self.tag_ids}

    def add_asset_to_album(self, album_id: str, asset_id: str) -> None:
        response = self.http.put(f"albums/{album_id}/assets", json={"ids": [asset_id]})
        self._check(response, "failed to add asset to album")

    def link_tag_to_asset(self, tag_id: str, asset_id: str) -> None:
        response = self.http.put(f"tags/{tag_id}/assets", json={"ids": [asset_id]})
        self._check(response, "failed to link tag to asset")

    def upload_image(self, path: Path, albums: list[str], tags: list[str]) -> str:
        asset = self.upload_asset(path)
        asset_id = asset["id"]
        for album in albums:
            try:
                self.add_asset_to_album(self.get_or_create_album(album), asset_id)
            except (httpx.HTTPError, RuntimeError, KeyError, ValueError) as error:
                LOG.error("Could not add %s to album %r: %s", path, album, error)
        try:
            tag_ids = self.create_tags(tags)
            for tag, tag_id in tag_ids.items():
                try:
                    self.link_tag_to_asset(tag_id, asset_id)
                except (httpx.HTTPError, RuntimeError) as error:
                    LOG.error("Could not apply tag %r to %s: %s", tag, path, error)
        except (httpx.HTTPError, RuntimeError, KeyError, ValueError) as error:
            LOG.error("Could not create tags for %s: %s", path, error)
        return asset.get("status") or "uploaded"


def process_image(path: Path, client: ImmichClient | None, dry_run: bool) -> bool:
    try:
        parsed = parse_filename(path)
        if parsed is None:
            raise ValueError(f"failed to parse filename: {path}")
        people, studio = parsed
        width, height = read_jpeg_dimensions(path)
        albums = [person.lower() for person in people]
        tags = list(albums)
        if studio:
            tags.append(studio.lower())
        tags.append(TAG_PORTRAIT if width <= height else TAG_LANDSCAPE)
        if dry_run:
            LOG.info("[DRY RUN] Would upload %s (albums=%s, tags=%s)", path, albums, tags)
            return True
        assert client is not None
        status = client.upload_image(path, albums, tags)
        LOG.info("Processed %s (status=%s)", path, status)
        return True
    except (OSError, ValueError, RuntimeError, httpx.HTTPError, KeyError) as error:
        LOG.error("Failed to process %s: %s", path, error)
        return False


def run(args: argparse.Namespace) -> int:
    logging.basicConfig(
        level=logging.DEBUG if args.debug else logging.INFO,
        format="%(levelname)s: %(message)s",
    )
    paths: list[Path] = []
    failures = 0
    for path in args.images:
        if not path.exists():
            LOG.error("Image does not exist, skipping: %s", path)
            failures += 1
        elif path.is_dir():
            LOG.warning("Skipping directory: %s", path)
        else:
            paths.append(path)
    LOG.info("Found %d images to process", len(paths))

    if args.move_uploaded_to:
        try:
            Path(args.move_uploaded_to).mkdir(parents=True, exist_ok=True)
        except OSError as error:
            LOG.error("Failed to create destination directory: %s", error)
            return 1

    client = None if args.dry_run else ImmichClient(args.server_url, args.api_key)
    try:
        with ThreadPoolExecutor(max_workers=5) as pool:
            futures = {pool.submit(process_image, path, client, args.dry_run): path for path in paths}
            for future in as_completed(futures):
                path = futures[future]
                try:
                    succeeded = future.result()
                except Exception:
                    LOG.exception("Unexpected error processing %s", path)
                    succeeded = False
                if not succeeded:
                    failures += 1
                    continue
                if args.move_uploaded_to:
                    destination = Path(args.move_uploaded_to) / path.name
                    if args.dry_run:
                        LOG.info("[DRY RUN] Would move %s to %s", path, destination.parent)
                    else:
                        try:
                            shutil.move(str(path), str(destination))
                            LOG.info("Moved %s to %s", path, destination)
                        except OSError as error:
                            LOG.error("Failed to move %s: %s", path, error)
                            failures += 1
    finally:
        if client is not None:
            client.close()
    LOG.info("Completed processing all images")
    return 1 if failures else 0


def main() -> int:
    return run(parse_args())


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