#!/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 collections import defaultdict
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:
        """Get existing album or create new one."""
        # Check cache first
        try:
            return self._album_id_by_name[name]
        except KeyError:
            pass

        # Try to find existing album
        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}")

        # Create new album
        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]:
        """Create or get tags, return mapping of tag_name -> tag_id."""
        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.info(f"Creating/retrieving 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:
        # Upload the image
        asset_response = client.upload_image(image_path)
        asset_id = asset_response["id"]
        status = asset_response.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

        # Handle albums
        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}")

        # Handle tags
        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_image_filename(image_path: Path) -> ParsedFilename | None:
    """Parse image filename to extract actor names."""
    # Remove extension and split by ' -- '
    filename = image_path.stem
    if " -- " in filename:
        actors_part = filename.split(" -- ")[0]
        # Split by comma and clean up whitespace
        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)

    logger.warning(f"Filename doesn't match expected format: {filename}")
    return None


def list_images(folder_path: Path) -> List[Path]:
    """List all image files in the folder."""
    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 group_images_by_actor(images: List[Path]) -> dict:
    """Group images by first actor name."""
    actor_groups = defaultdict(list)

    for image in images:
        parsed = parse_image_filename(image)
        if parsed:
            first_actor = parsed.actors[0].lower()
            actor_groups[first_actor].append(image)
            logger.debug(f"Assigned {image.name} to actor: {first_actor}")
        else:
            logger.warning(f"Skipping image with no actors: {image.name}")

    logger.info(f"Grouped images into {len(actor_groups)} actor categories")
    return actor_groups


def parse_args():
    """Parse command-line arguments."""
    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

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

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

    # Create _uploaded directory
    # if they're all in the same directory
    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:
        try:
            # Parse filename to extract actor names
            parsed = parse_image_filename(image_path)
            if not parsed:
                logger.warning(f"Skipping image with unparseable filename: {image_path.name}")
                continue

            # Convert actor names to lowercase for albums and tags

            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)  # Use rename to move the file
        except Exception as e:
            logger.error(f"Error processing image {image_path.name}: {e}")


if __name__ == "__main__":
    main()