"""
Asset management system for pyPhotoAlbum with automatic reference counting
"""

import hashlib
import os
import shutil
from typing import Dict, List, Optional, Tuple
from pathlib import Path


def compute_file_md5(file_path: str) -> Optional[str]:
    """
    Compute MD5 hash of a file.

    Args:
        file_path: Path to the file

    Returns:
        MD5 hash as hex string, or None if file doesn't exist
    """
    if not os.path.exists(file_path):
        return None

    hash_md5 = hashlib.md5()
    try:
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                hash_md5.update(chunk)
        return hash_md5.hexdigest()
    except Exception as e:
        print(f"AssetManager: Error computing MD5 for {file_path}: {e}")
        return None


class AssetManager:
    """Manages project assets with automatic reference counting and cleanup"""

    def __init__(self, project_folder: str):
        """
        Initialize AssetManager.

        Args:
            project_folder: Root folder for the project
        """
        self.project_folder = project_folder
        self.assets_folder = os.path.join(project_folder, "assets")
        self.reference_counts: Dict[str, int] = {}  # {relative_path: count}
        self.asset_hashes: Dict[str, str] = {}  # {relative_path: md5_hash}

        # Create assets folder if it doesn't exist
        os.makedirs(self.assets_folder, exist_ok=True)

    def import_asset(self, source_path: str) -> str:
        """
        Import an asset into the project by copying it to the assets folder.

        Args:
            source_path: Path to the source file

        Returns:
            Relative path to the imported asset (e.g., "assets/photo_001.jpg")
        """
        if not os.path.exists(source_path):
            raise FileNotFoundError(f"Source file not found: {source_path}")

        # Get filename and extension
        filename = os.path.basename(source_path)
        name, ext = os.path.splitext(filename)

        # Find a unique filename if there's a collision
        counter = 1
        dest_filename = filename
        dest_path = os.path.join(self.assets_folder, dest_filename)

        while os.path.exists(dest_path):
            dest_filename = f"{name}_{counter:03d}{ext}"
            dest_path = os.path.join(self.assets_folder, dest_filename)
            counter += 1

        # Copy the file
        shutil.copy2(source_path, dest_path)

        # Get relative path from project folder (for storage/serialization)
        relative_path = os.path.relpath(dest_path, self.project_folder)

        # Initialize reference count
        self.reference_counts[relative_path] = 1

        print(f"AssetManager: Imported {source_path} → {dest_path} (stored as {relative_path}, refs=1)")

        # Return relative path for storage in elements
        return relative_path

    def acquire_reference(self, asset_path: str):
        """
        Increment the reference count for an asset.

        Args:
            asset_path: Relative path to the asset
        """
        if not asset_path:
            return

        if asset_path in self.reference_counts:
            self.reference_counts[asset_path] += 1
            print(f"AssetManager: Acquired reference to {asset_path} (refs={self.reference_counts[asset_path]})")
        else:
            # Asset might exist from a loaded project
            full_path = os.path.join(self.project_folder, asset_path)
            if os.path.exists(full_path):
                self.reference_counts[asset_path] = 1
                print(f"AssetManager: Acquired reference to existing asset {asset_path} (refs=1)")
            else:
                print(f"AssetManager: Warning - asset not found: {asset_path}")

    def release_reference(self, asset_path: str):
        """
        Decrement the reference count for an asset.
        If count reaches zero, delete the asset file.

        Args:
            asset_path: Relative path to the asset
        """
        if not asset_path:
            return

        if asset_path not in self.reference_counts:
            print(f"AssetManager: Warning - attempting to release unknown asset: {asset_path}")
            return

        self.reference_counts[asset_path] -= 1
        print(f"AssetManager: Released reference to {asset_path} (refs={self.reference_counts[asset_path]})")

        if self.reference_counts[asset_path] <= 0:
            # No more references - safe to delete
            full_path = os.path.join(self.project_folder, asset_path)
            try:
                if os.path.exists(full_path):
                    os.remove(full_path)
                    print(f"AssetManager: Deleted unused asset {asset_path}")
                del self.reference_counts[asset_path]
            except Exception as e:
                print(f"AssetManager: Error deleting asset {asset_path}: {e}")

    def get_absolute_path(self, relative_path: str) -> str:
        """
        Convert a relative asset path to an absolute path.

        Args:
            relative_path: Relative path from project folder

        Returns:
            Absolute path to the asset
        """
        return os.path.join(self.project_folder, relative_path)

    def get_reference_count(self, asset_path: str) -> int:
        """
        Get the current reference count for an asset.

        Args:
            asset_path: Relative path to the asset

        Returns:
            Reference count (0 if not tracked)
        """
        return self.reference_counts.get(asset_path, 0)

    def serialize(self) -> Dict:
        """Serialize asset manager state"""
        return {
            "reference_counts": self.reference_counts,
            "asset_hashes": self.asset_hashes,
        }

    def deserialize(self, data: Dict):
        """Deserialize asset manager state"""
        self.reference_counts = data.get("reference_counts", {})
        self.asset_hashes = data.get("asset_hashes", {})
        print(f"AssetManager: Loaded {len(self.reference_counts)} asset references")

    def compute_asset_hash(self, asset_path: str) -> Optional[str]:
        """
        Compute and cache the MD5 hash for an asset.

        Args:
            asset_path: Relative path to the asset

        Returns:
            MD5 hash as hex string, or None if computation fails
        """
        full_path = self.get_absolute_path(asset_path)
        md5_hash = compute_file_md5(full_path)
        if md5_hash:
            self.asset_hashes[asset_path] = md5_hash
        return md5_hash

    def compute_all_hashes(self) -> Dict[str, str]:
        """
        Compute MD5 hashes for all assets in the assets folder.

        Returns:
            Dictionary mapping relative paths to MD5 hashes
        """
        self.asset_hashes.clear()

        if not os.path.exists(self.assets_folder):
            return self.asset_hashes

        for root, dirs, files in os.walk(self.assets_folder):
            for filename in files:
                file_path = os.path.join(root, filename)
                relative_path = os.path.relpath(file_path, self.project_folder)
                md5_hash = compute_file_md5(file_path)
                if md5_hash:
                    self.asset_hashes[relative_path] = md5_hash

        print(f"AssetManager: Computed hashes for {len(self.asset_hashes)} assets")
        return self.asset_hashes

    def find_duplicates(self) -> Dict[str, List[str]]:
        """
        Find duplicate assets based on MD5 hash.

        Returns:
            Dictionary mapping MD5 hash to list of asset paths with that hash.
            Only includes hashes with more than one file.
        """
        # Compute hashes if not already done
        if not self.asset_hashes:
            self.compute_all_hashes()

        # Group assets by hash
        hash_to_paths: Dict[str, List[str]] = {}
        for path, md5_hash in self.asset_hashes.items():
            if md5_hash not in hash_to_paths:
                hash_to_paths[md5_hash] = []
            hash_to_paths[md5_hash].append(path)

        # Filter to only duplicates (more than one file with same hash)
        duplicates = {h: paths for h, paths in hash_to_paths.items() if len(paths) > 1}

        if duplicates:
            total_dups = sum(len(paths) - 1 for paths in duplicates.values())
            print(f"AssetManager: Found {total_dups} duplicate files in {len(duplicates)} groups")

        return duplicates

    def deduplicate_assets(self, update_references_callback=None) -> Tuple[int, int]:
        """
        Remove duplicate assets, keeping one canonical copy and updating references.

        Args:
            update_references_callback: Optional callback function that takes
                (old_path, new_path) to update external references (e.g., ImageData elements)

        Returns:
            Tuple of (files_removed, bytes_saved)
        """
        duplicates = self.find_duplicates()
        if not duplicates:
            print("AssetManager: No duplicates found")
            return (0, 0)

        files_removed = 0
        bytes_saved = 0

        for md5_hash, paths in duplicates.items():
            # Sort paths to get consistent canonical path (first alphabetically)
            paths.sort()
            canonical_path = paths[0]

            # Remove duplicates and update references
            for dup_path in paths[1:]:
                full_dup_path = self.get_absolute_path(dup_path)

                # Get file size before deletion
                try:
                    file_size = os.path.getsize(full_dup_path)
                except OSError:
                    file_size = 0

                # Update references if callback provided
                if update_references_callback:
                    update_references_callback(dup_path, canonical_path)

                # Transfer reference count to canonical path
                if dup_path in self.reference_counts:
                    dup_refs = self.reference_counts[dup_path]
                    if canonical_path in self.reference_counts:
                        self.reference_counts[canonical_path] += dup_refs
                    else:
                        self.reference_counts[canonical_path] = dup_refs
                    del self.reference_counts[dup_path]

                # Delete the duplicate file
                try:
                    if os.path.exists(full_dup_path):
                        os.remove(full_dup_path)
                        files_removed += 1
                        bytes_saved += file_size
                        print(f"AssetManager: Removed duplicate {dup_path} (kept {canonical_path})")
                except Exception as e:
                    print(f"AssetManager: Error removing duplicate {dup_path}: {e}")

                # Remove from hash tracking
                if dup_path in self.asset_hashes:
                    del self.asset_hashes[dup_path]

        print(f"AssetManager: Deduplication complete - removed {files_removed} files, saved {bytes_saved} bytes")
        return (files_removed, bytes_saved)

    def get_duplicate_stats(self) -> Tuple[int, int, int]:
        """
        Get statistics about duplicate assets without modifying anything.

        Returns:
            Tuple of (duplicate_groups, total_duplicate_files, estimated_bytes_to_save)
        """
        duplicates = self.find_duplicates()
        if not duplicates:
            return (0, 0, 0)

        duplicate_groups = len(duplicates)
        total_duplicate_files = sum(len(paths) - 1 for paths in duplicates.values())

        # Calculate bytes that would be saved
        bytes_to_save = 0
        for paths in duplicates.values():
            # Skip the first (canonical) file, count size of the rest
            for dup_path in paths[1:]:
                full_path = self.get_absolute_path(dup_path)
                try:
                    bytes_to_save += os.path.getsize(full_path)
                except OSError:
                    pass

        return (duplicate_groups, total_duplicate_files, bytes_to_save)

    def find_unused_assets(self) -> List[str]:
        """
        Find assets that exist in the assets folder but have no references.

        Returns:
            List of relative paths to unused assets
        """
        unused: list[str] = []

        if not os.path.exists(self.assets_folder):
            return unused

        # Get all files in assets folder
        for root, dirs, files in os.walk(self.assets_folder):
            for filename in files:
                file_path = os.path.join(root, filename)
                relative_path = os.path.relpath(file_path, self.project_folder)

                # Check if this asset has any references
                ref_count = self.reference_counts.get(relative_path, 0)
                if ref_count <= 0:
                    unused.append(relative_path)

        if unused:
            print(f"AssetManager: Found {len(unused)} unused assets")

        return unused

    def get_unused_stats(self) -> Tuple[int, int]:
        """
        Get statistics about unused assets without modifying anything.

        Returns:
            Tuple of (unused_file_count, total_bytes)
        """
        unused = self.find_unused_assets()
        if not unused:
            return (0, 0)

        total_bytes = 0
        for asset_path in unused:
            full_path = self.get_absolute_path(asset_path)
            try:
                total_bytes += os.path.getsize(full_path)
            except OSError:
                pass

        return (len(unused), total_bytes)

    def remove_unused_assets(self) -> Tuple[int, int]:
        """
        Remove all unused assets from the assets folder.

        Returns:
            Tuple of (files_removed, bytes_freed)
        """
        unused = self.find_unused_assets()
        if not unused:
            print("AssetManager: No unused assets to remove")
            return (0, 0)

        files_removed = 0
        bytes_freed = 0

        for asset_path in unused:
            full_path = self.get_absolute_path(asset_path)

            try:
                file_size = os.path.getsize(full_path)
            except OSError:
                file_size = 0

            try:
                if os.path.exists(full_path):
                    os.remove(full_path)
                    files_removed += 1
                    bytes_freed += file_size
                    print(f"AssetManager: Removed unused asset {asset_path}")

                # Clean up tracking
                if asset_path in self.reference_counts:
                    del self.reference_counts[asset_path]
                if asset_path in self.asset_hashes:
                    del self.asset_hashes[asset_path]

            except Exception as e:
                print(f"AssetManager: Error removing unused asset {asset_path}: {e}")

        print(f"AssetManager: Removed {files_removed} unused assets, freed {bytes_freed} bytes")
        return (files_removed, bytes_freed)
