#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["httpx"]
# ///
import dataclasses
import os
import re
import logging
import argparse
import uuid
import hashlib
from pathlib import Path
from typing import List, Dict
from datetime import datetime
import httpx

# Configure logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
logging.getLogger("httpx").setLevel(logging.WARNING)

IMMICH_SERVER_URL = "https://ph.abdus.dev"  # Replace with your Immich server URL
IMMICH_API_KEY = os.getenv("IMMICH_API_KEY")


class ImmichClient:
    """Client for interacting with Immich API."""

    def __init__(self, base_url: str, api_key: str):
        self.base_url = base_url.rstrip("/") + "/api"
        self.api_key = api_key
        self.device_uuid = str(uuid.uuid4())
        self.client = httpx.Client(
            timeout=60,
            headers={
                "x-api-key": self.api_key,
            },
            base_url=self.base_url,
        )
        self._album_id_by_name = {}
        self._tag_id_by_name = {}

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.client.close()

    def get_or_create_album(self, name: str) -> str:
        try:
            return self._album_id_by_name[name]
        except KeyError:
            pass

        try:
            res = self.client.get("/albums")
            res.raise_for_status()
            albums = res.json()

            for album in albums:
                existing_album = album["albumName"]
                self._album_id_by_name[existing_album] = album["id"]
                if existing_album == name:
                    logger.info(f"Found existing album: {name} (ID: {album['id']})")
                    return album["id"]
        except Exception as e:
            logger.warning(f"Error searching for album: {e}")

        logger.info(f"Creating new album: {name}")
        res = self.client.post("/albums", json={"albumName": name})
        res.raise_for_status()
        album = res.json()
        album_id = album["id"]
        self._album_id_by_name[name] = album_id
        return album_id

    def create_tags(self, tags: list[str]) -> dict[str, str]:
        if not tags:
            return {}

        # Check cache for existing tags
        cached_tags = {}
        missing_tags = []

        for tag in tags:
            if tag in self._tag_id_by_name:
                cached_tags[tag] = self._tag_id_by_name[tag]
            else:
                missing_tags.append(tag)

        if not missing_tags:
            logger.debug(f"All tags found in cache: {tags}")
            return cached_tags

        logger.debug(f"Creating tags: {missing_tags}")
        res = self.client.put("/tags", json={"tags": missing_tags})
        res.raise_for_status()

        # Update cache with new tags
        new_tags = {}
        for tag in res.json():
            tag_name = tag["name"]
            tag_id = tag["id"]
            self._tag_id_by_name[tag_name] = tag_id
            new_tags[tag_name] = tag_id

        # Combine cached and new tags
        all_tags = {**cached_tags, **new_tags}
        return all_tags

    def add_asset_to_album(self, album_id: str, asset_id: str) -> None:
        """Add asset to album."""
        logger.debug(f"Adding asset {asset_id} to album {album_id}")
        res = self.client.put(f"/albums/{album_id}/assets", json={"ids": [asset_id]})
        res.raise_for_status()

    def link_tag_to_asset(self, tag_id: str, asset_id: str) -> None:
        """Link tag to asset."""
        logger.debug(f"Linking tag {tag_id} to asset {asset_id}")
        res = self.client.put(f"/tags/{tag_id}/assets", json={"ids": [asset_id]})
        res.raise_for_status()

    def upload_image(self, image_path: Path) -> Dict:
        """Upload single image to Immich."""
        logger.info(f"Uploading image: {image_path}")

        stat = image_path.stat()
        filename = image_path.name
        ext = image_path.suffix.lower()

        device_asset_id = f"{filename}-{stat.st_size}"

        file_time = datetime.fromtimestamp(stat.st_mtime).isoformat() + "Z"

        # Calculate SHA1 checksum for duplicate detection
        sha1_hash = hashlib.sha1()
        with image_path.open("rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                sha1_hash.update(chunk)
        checksum = sha1_hash.hexdigest()

        with image_path.open("rb") as f:
            response = self.client.post(
                "/assets",
                headers={
                    "x-immich-checksum": checksum,
                },
                files={
                    "assetData": (filename, f, "image/jpeg"),
                },
                data={
                    "deviceAssetId": device_asset_id,
                    "deviceId": self.device_uuid,
                    "assetType": "image",
                    "fileCreatedAt": file_time,
                    "fileModifiedAt": file_time,
                    "isFavorite": "false",
                    "fileExtension": ext,
                    "duration": "0",
                    "isReadOnly": "false",
                    "isArchived": "false",
                },
            )
            response.raise_for_status()
            return response.json()


def upload_image(
    client: ImmichClient,
    image_path: Path,
    albums: list[str] | None = None,
    tags: list[str] | None = None,
) -> None:
    """
    Upload a single image to Immich with optional albums and tags.

    Args:
        image_path: Path to the image file
        albums: List of album names to add the image to
        tags: List of tags to apply to the image
        immich_client: ImmichClient instance to use (defaults to global instance)
    """

    albums = albums or []
    tags = tags or []

    try:
        asset = client.upload_image(image_path)
        asset_id = asset["id"]
        status = asset.get("status", "uploaded")

        logger.info(f"Uploaded: {image_path.name} (status={status})")

        # Skip album/tag processing for duplicates
        # if status == "duplicate":
        #    logger.info(f"Skipping album/tag processing for duplicate: {image_path.name}")
        #    return

        album_ids = []
        for album_name in albums:
            try:
                album_id = client.get_or_create_album(album_name)
                album_ids.append(album_id)
                client.add_asset_to_album(album_id, asset_id)
                logger.debug(f"Added to album: {album_name}")
            except Exception as e:
                logger.error(f"Failed to add to album {album_name}: {e}")

        if tags:
            try:
                tag_id_by_name = client.create_tags(tags)
                for tag_name, tag_id in tag_id_by_name.items():
                    client.link_tag_to_asset(tag_id, asset_id)
                    logger.debug(f"Applied tag: {tag_name}")
            except Exception as e:
                logger.error(f"Failed to apply tags: {e}")

    except Exception as e:
        logger.error(f"Failed to upload {image_path}: {e}")
        raise


@dataclasses.dataclass
class ParsedFilename:
    actors: list[str]
    studio: str | None = None


def parse_filename(image_path: Path) -> ParsedFilename | None:
    filename = image_path.stem
    if " -- " in filename:
        actors_part = filename.split(" -- ")[0]
        actors = [actor.strip() for actor in actors_part.split(",")]
        studio: str | None = None
        if m := re.search(r" -- @(\S+)", filename):
            studio = m.group(1)
        return ParsedFilename(actors=actors, studio=studio)

    return None


def list_images(folder_path: Path) -> List[Path]:
    image_extensions = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".tiff", ".webp"}
    images = []

    for ext in image_extensions:
        images.extend(folder_path.glob(f"*{ext}"))

    return images


def parse_args():
    parser = argparse.ArgumentParser(description="Upload images to Immich, grouping them by actor name from filename")
    parser.add_argument("image_paths", type=Path, nargs="+", help="Path to the images to upload")
    return parser.parse_args()


def main():
    args = parse_args()
    image_paths: list[Path] = args.image_paths

    if not image_paths:
        logger.error("No images found to process")
        return

    logger.info(f"Found {len(image_paths)} images to process")

    if not all(image_paths[0].parent == path.parent for path in image_paths):
        logger.error("All images must be in the same directory")
        return
    uploaded_dir = image_paths[0].parent / "_uploaded"
    uploaded_dir.mkdir(exist_ok=True)

    immich = ImmichClient(IMMICH_SERVER_URL, IMMICH_API_KEY)

    for image_path in image_paths:
        if not image_path.is_file():
            logger.warning(f"Skipping non-file path: {image_path}")
            continue

        try:
            parsed = parse_filename(image_path)
            if not parsed:
                logger.warning(f"Skipping image with unparseable filename: {image_path.name}")
                continue

            logger.info(f"Processing {image_path.name} with actors: {', '.join(parsed.actors)}")

            actor_names = [actor.lower() for actor in parsed.actors]
            tags = [*actor_names]
            if parsed.studio:
                tags.append(parsed.studio.lower())

            upload_image(
                client=immich,
                image_path=image_path,
                albums=actor_names,
                tags=tags,
            )

            dest_path = uploaded_dir / image_path.name
            logger.info(f"Moving {image_path.name} to {dest_path}")
            image_path.rename(dest_path)
        except Exception as e:
            logger.error(f"Error processing image {image_path.name}: {e}")


if __name__ == "__main__":
    main()