#!/usr/bin/env python3
"""
iPod Device Handler Module for iPod Nano Transfer Tool
Handles device detection, mounting, and file transfer to iPod Nano
"""

import os
import sys
import json
import logging
import platform
import shutil
import subprocess
from typing import List, Dict, Optional, Tuple, Any
from pathlib import Path

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class IPodDevice:
    """Class to handle iPod device detection and file transfer"""
    
    def __init__(self, mount_point: Optional[str] = None):
        """
        Initialize the iPod device handler
        
        Args:
            mount_point: Path to iPod mount point (optional)
        """
        self.mount_point = mount_point
        self.system = platform.system()
        self.device_info = {}
        self._block_device = None  # /dev/sdd1 etc.
        self._parent_drive = None  # /dev/sdd (for power-off)
        self.logger = logger
    
    def detect_devices(self) -> List[Dict[str, str]]:
        """
        Detect connected iPod devices
        
        Returns:
            List of dictionaries with device information
        """
        devices = []
        system = platform.system()
        
        try:
            if system == "Linux":
                devices = self._detect_linux()
            elif system == "Darwin":
                devices = self._detect_macos()
            elif system == "Windows":
                devices = self._detect_windows()
        except Exception as e:
            self.logger.error(f"Error detecting devices: {e}")
            
        # If no devices found, check for manual mount point
        if not devices and self.mount_point and os.path.exists(self.mount_point):
            devices.append({
                "id": "manual",
                "name": os.path.basename(self.mount_point),
                "model": "Unknown",
                "mount_point": self.mount_point,
                "mounted": True,
            })
            
        return devices
    
    def _detect_linux(self) -> List[Dict[str, str]]:
        """Detect iPod on Linux using lsblk and udisksctl"""
        devices = []
        
        # Method 1: lsblk JSON output - find devices with IPOD label
        try:
            result = subprocess.run(
                ["lsblk", "-o", "NAME,LABEL,MOUNTPOINT,TYPE", "-J"],
                capture_output=True, check=False, text=True
            )
            
            if result.returncode == 0:
                data = json.loads(result.stdout)
                for device in data.get("blockdevices", []):
                    if device.get("type") != "disk":
                        continue
                    parent_name = device.get("name", "")
                    for child in device.get("children", []):
                        label = child.get("label", "")
                        mountpoint = child.get("mountpoint")
                        child_name = child.get("name", "")
                        
                        if label and "ipod" in label.lower():
                            dev_path = f"/dev/{child_name}"
                            parent_dev = f"/dev/{parent_name}"
                            
                            device_entry = {
                                "id": dev_path,
                                "name": label,
                                "model": "iPod Nano 7",
                                "mount_point": mountpoint,
                                "mounted": mountpoint is not None,
                                "_parent_drive": parent_dev,
                            }
                            devices.append(device_entry)
        except (FileNotFoundError, json.JSONDecodeError) as e:
            self.logger.warning(f"lsblk JSON failed: {e}")
        
        # Method 2: udisksctl dump - find iPod drives
        if not devices:
            try:
                result = subprocess.run(
                    ["udisksctl", "dump"],
                    capture_output=True, check=False, text=True
                )
                
                if result.returncode == 0:
                    current_block = None
                    current_drive = None
                    current_label = None
                    current_mount = None
                    
                    for line in result.stdout.splitlines():
                        if "/org/freedesktop/UDisks2/block_devices/" in line:
                            # Save previous if it was an iPod
                            if current_block and current_label and "ipod" in current_label.lower():
                                devices.append({
                                    "id": current_block,
                                    "name": current_label,
                                    "model": "iPod",
                                    "mount_point": current_mount if current_mount else None,
                                    "mounted": current_mount is not None,
                                    "_parent_drive": self._get_parent_drive_from_udisks(current_block),
                                })
                            current_block = line.split("/block_devices/")[-1].strip().rstrip(":")
                            current_drive = None
                            current_label = None
                            current_mount = None
                        
                        if "IdLabel:" in line and current_block:
                            val = line.split(":", 1)[1].strip()
                            if val and val != "''":
                                current_label = val.strip("'")
                        if "MountPoints:" in line and current_block:
                            val = line.split(":", 1)[1].strip()
                            if val and val != "''" and val != "b''":
                                # MountPoints is a byte array like b'/run/media/mat/IPOD'
                                import ast
                                try:
                                    decoded = ast.literal_eval(val)
                                    if isinstance(decoded, bytes):
                                        current_mount = decoded.decode("utf-8", errors="replace")
                                    elif isinstance(decoded, str):
                                        current_mount = decoded
                                except (ValueError, SyntaxError):
                                    current_mount = val.strip("'\"")
                    
                    # Don't forget the last entry
                    if current_block and current_label and "ipod" in current_label.lower():
                        devices.append({
                            "id": current_block,
                            "name": current_label,
                            "model": "iPod",
                            "mount_point": current_mount if current_mount else None,
                            "mounted": current_mount is not None,
                            "_parent_drive": self._get_parent_drive_from_udisks(current_block),
                        })
            except FileNotFoundError:
                self.logger.warning("udisksctl not found")
        
        # Method 3: Check common mount directories for already-mounted iPods
        for device in devices:
            if not device["mounted"]:
                for base in ["/run/media", "/media", "/mnt"]:
                    if os.path.exists(base):
                        for user_or_device in os.listdir(base):
                            user_path = os.path.join(base, user_or_device)
                            if os.path.isdir(user_path):
                                # Could be /run/media/mat/IPOD or /media/IPOD
                                check_paths = [user_path]
                                if os.path.isdir(user_path):
                                    for sub in os.listdir(user_path):
                                        sub_path = os.path.join(user_path, sub)
                                        if "ipod" in sub.lower() and os.path.isdir(sub_path):
                                            check_paths.append(sub_path)
                                
                                for check_path in check_paths:
                                    if os.path.exists(os.path.join(check_path, "iPod_Control")):
                                        device["mount_point"] = check_path
                                        device["mounted"] = True
                                        break
                            if device["mounted"]:
                                break
                    if device["mounted"]:
                        break
        
        return devices
    
    def _get_parent_drive_from_udisks(self, block_name: str) -> Optional[str]:
        """Get parent drive device from udisksctl for a partition"""
        try:
            # /dev/sdd1 -> /dev/sdd
            if block_name and block_name.startswith("/dev/"):
                name = block_name[5:]  # sdd1
                # Strip trailing partition number
                base = name.rstrip("0123456789")
                if base:
                    return f"/dev/{base}"
        except Exception:
            pass
        return None
    
    def _detect_macos(self) -> List[Dict[str, str]]:
        """Detect iPod on macOS"""
        devices = []
        
        # Check /Volumes for iPod
        volumes_dir = "/Volumes"
        if os.path.exists(volumes_dir):
            for volume in os.listdir(volumes_dir):
                volume_path = os.path.join(volumes_dir, volume)
                if "iPod" in volume and os.path.isdir(volume_path):
                    if os.path.exists(os.path.join(volume_path, "iPod_Control")):
                        devices.append({
                            "id": "unknown",
                            "name": volume,
                            "model": "Unknown",
                            "mount_point": volume_path,
                            "mounted": True,
                        })
        
        return devices
    
    def _detect_windows(self) -> List[Dict[str, str]]:
        """Detect iPod on Windows"""
        devices = []
        import string
        for letter in string.ascii_uppercase:
            drive = f"{letter}:\\"
            if os.path.exists(drive):
                if os.path.exists(os.path.join(drive, "iPod_Control")):
                    devices.append({
                        "id": letter,
                        "name": "iPod",
                        "model": "Unknown",
                        "mount_point": drive,
                        "mounted": True,
                    })
        return devices
    
    def mount_device(self, device_id: Optional[str] = None, detected_mount: Optional[str] = None) -> Optional[str]:
        """
        Mount iPod device using udisksctl (Linux) or verify existing mount
        
        Args:
            device_id: Device path like /dev/sdd1
            detected_mount: Already detected mount point
            
        Returns:
            Mount point if successful, None otherwise
        """
        # If device is already mounted at detected path, just use it
        if detected_mount and os.path.exists(detected_mount):
            if os.path.exists(os.path.join(detected_mount, "iPod_Control")):
                self.mount_point = detected_mount
                self.logger.info(f"Using already mounted device at: {detected_mount}")
                return detected_mount
        
        if self.mount_point and os.path.exists(self.mount_point):
            self.logger.info(f"Device already mounted at: {self.mount_point}")
            return self.mount_point
        
        try:
            if self.system == "Linux":
                return self._mount_linux(device_id)
            elif self.system == "Darwin":
                return self._mount_macos()
            elif self.system == "Windows":
                return self._mount_windows()
        except subprocess.SubprocessError as e:
            self.logger.error(f"Failed to mount device: {e}")
            return None
        
        return None
    
    def _mount_linux(self, device_id: Optional[str]) -> Optional[str]:
        """Mount on Linux using udisksctl"""
        # device_id should be like /dev/sdd1
        if not device_id or not device_id.startswith("/dev/"):
            self.logger.error("No valid device path provided for mounting")
            return None
        
        if not os.path.exists(device_id):
            self.logger.error(f"Device {device_id} does not exist")
            return None
        
        try:
            result = subprocess.run(
                ["udisksctl", "mount", "-b", device_id],
                capture_output=True, check=False, text=True
            )
            
            if result.returncode != 0:
                self.logger.error(f"udisksctl mount failed: {result.stderr}")
                return None
            
            # Parse mount point from output: "Mounted /dev/sdd1 at /run/media/mat/IPOD.\n"
            output = result.stdout.strip()
            mount_point = None
            for line in output.splitlines():
                if "Mounted" in line and "at" in line:
                    # "Mounted /dev/sdd1 at /run/media/mat/IPOD."
                    parts = line.split(" at ")
                    if len(parts) >= 2:
                        mp = parts[1].strip().rstrip(".")
                        if os.path.exists(mp):
                            mount_point = mp
            
            if not mount_point:
                # Fallback: check /run/media for the device
                for user_dir in os.listdir("/run/media"):
                    user_path = os.path.join("/run/media", user_dir)
                    if os.path.isdir(user_path):
                        for d in os.listdir(user_path):
                            d_path = os.path.join(user_path, d)
                            if os.path.exists(os.path.join(d_path, "iPod_Control")):
                                mount_point = d_path
                                break
                    if mount_point:
                        break
            
            if mount_point:
                self.mount_point = mount_point
                self._block_device = device_id
                # Extract parent drive: /dev/sdd1 -> /dev/sdd
                name = device_id[5:]
                self._parent_drive = f"/dev/{name.rstrip('0123456789')}"
                self.logger.info(f"Device mounted at: {mount_point}")
                return mount_point
            
            self.logger.error("Could not determine mount point")
            return None
            
        except FileNotFoundError:
            self.logger.error("udisksctl not found. Please install udisks2.")
            return None
    
    def _mount_macos(self) -> Optional[str]:
        """Mount on macOS - devices auto-mount, just find the mount point"""
        result = subprocess.run(
            ["diskutil", "list"],
            capture_output=True, check=False, text=True
        )
        
        if result.returncode == 0:
            for line in result.stdout.splitlines():
                if "iPod" in line:
                    df_result = subprocess.run(
                        ["df", "-h"],
                        capture_output=True, check=False, text=True
                    )
                    for df_line in df_result.stdout.splitlines():
                        if "iPod" in df_line:
                            parts = df_line.split()
                            if len(parts) >= 9:
                                mount_point = parts[8]
                                self.mount_point = mount_point
                                return mount_point
        
        return None
    
    def _mount_windows(self) -> Optional[str]:
        """Mount on Windows"""
        self.logger.warning("Windows mounting not fully implemented")
        return None
    
    def unmount_device(self) -> bool:
        """
        Unmount and safely eject iPod device
        
        Returns:
            True if successful, False otherwise
        """
        if not self.mount_point and not self._block_device:
            self.logger.warning("No device mounted")
            return False
        
        try:
            if self.system == "Linux":
                return self._unmount_linux()
            elif self.system == "Darwin":
                return self._unmount_macos()
            elif self.system == "Windows":
                return self._unmount_windows()
        except subprocess.SubprocessError as e:
            self.logger.error(f"Failed to unmount device: {e}")
            return False
        
        return False
    
    def _unmount_linux(self) -> bool:
        """Unmount and power-off on Linux using udisksctl"""
        success = True
        
        # Step 1: Unmount
        if self.mount_point or self._block_device:
            device = self._block_device
            if not device:
                # Try to find the device from mount point
                result = subprocess.run(
                    ["findmnt", "-n", "-o", "SOURCE", self.mount_point],
                    capture_output=True, check=False, text=True
                )
                if result.returncode == 0 and result.stdout.strip():
                    device = result.stdout.strip()
            
            if device:
                result = subprocess.run(
                    ["udisksctl", "unmount", "-b", device],
                    capture_output=True, check=False, text=True
                )
                
                if result.returncode == 0:
                    self.logger.info(f"Device unmounted: {device}")
                else:
                    self.logger.warning(f"Unmount warning: {result.stderr}")
                    # Not a fatal error - device might already be unmounted
        
        # Step 2: Power off the drive (safely remove)
        drive = self._parent_drive
        if not drive and self._block_device:
            name = self._block_device[5:]
            drive = f"/dev/{name.rstrip('0123456789')}"
        
        if drive:
            result = subprocess.run(
                ["udisksctl", "power-off", "-b", drive],
                capture_output=True, check=False, text=True
            )
            
            if result.returncode == 0:
                self.logger.info(f"Device powered off: {drive}")
            else:
                self.logger.warning(f"Power-off warning: {result.stderr}")
                # Not fatal - USB may have been unplugged already
        
        self.mount_point = None
        self._block_device = None
        self._parent_drive = None
        return success
    
    def _unmount_macos(self) -> bool:
        """Eject on macOS"""
        if self.mount_point:
            subprocess.run(["diskutil", "eject", self.mount_point], check=True)
            self.logger.info(f"Device ejected: {self.mount_point}")
            self.mount_point = None
            return True
        return False
    
    def _unmount_windows(self) -> bool:
        """Unmount on Windows"""
        self.logger.warning("Windows unmounting not fully implemented")
        return False
    
    def is_mounted(self) -> bool:
        """Check if iPod is currently mounted"""
        if self.mount_point and os.path.exists(self.mount_point):
            return os.path.exists(os.path.join(self.mount_point, "iPod_Control"))
        return False
    
    def get_device_info(self) -> Dict[str, Any]:
        """
        Get information about the mounted device
        
        Returns:
            Dictionary with device information
        """
        if not self.mount_point or not os.path.exists(self.mount_point):
            self.logger.warning("No device mounted")
            return {}
            
        info = {
            "mount_point": self.mount_point,
            "total_space": 0,
            "free_space": 0,
            "used_space": 0
        }
        
        try:
            # Get disk usage
            usage = shutil.disk_usage(self.mount_point)
            info["total_space"] = usage.total
            info["free_space"] = usage.free
            info["used_space"] = usage.used
            
            # Try to get track count from Nano7 database
            try:
                from ipod_nano7_db import Nano7Database
                nano7_db = Nano7Database(self.mount_point)
                info["track_count"] = nano7_db.get_track_count()
                info["has_nano7_db"] = True
            except Exception:
                info["has_nano7_db"] = False
            
            return info
            
        except Exception as e:
            self.logger.error(f"Failed to get device info: {e}")
            return info
    
    def transfer_file(self, source_file: str, relative_path: Optional[str] = None) -> bool:
        """
        Transfer file to iPod device
        
        Args:
            source_file: Path to source file
            relative_path: Relative path on iPod (optional)
            
        Returns:
            True if successful, False otherwise
        """
        if not self.mount_point or not os.path.exists(self.mount_point):
            self.logger.warning("No device mounted")
            return False
            
        if not os.path.exists(source_file):
            self.logger.warning(f"Source file does not exist: {source_file}")
            return False
            
        try:
            # Try to use Nano7 database for proper track registration
            try:
                from ipod_nano7_db import Nano7Database
                nano7_db = Nano7Database(self.mount_point)
                result = nano7_db.add_track(source_file)
                if result is None:
                    # Track already exists (duplicate)
                    self.logger.info(f"Track already on iPod, skipping: {os.path.basename(source_file)}")
                    return True  # Not an error - track is already there
                self.logger.info(f"Track added to iPod database: {result['title']}")
                return True
            except ImportError:
                self.logger.warning("ipod_nano7_db not available, using simple file copy")
            except Exception as e:
                self.logger.error(f"Nano7 database transfer failed: {e}, falling back to simple copy")
            
            # Fallback: simple file copy (won't show in iPod interface)
            if relative_path:
                dest_dir = os.path.join(self.mount_point, relative_path)
            else:
                dest_dir = os.path.join(self.mount_point, "Music")
                parts = Path(source_file).parts
                if len(parts) >= 3:
                    artist = parts[-3]
                    album = parts[-2]
                    dest_dir = os.path.join(self.mount_point, "Music", artist, album)
            
            os.makedirs(dest_dir, exist_ok=True)
            filename = os.path.basename(source_file)
            dest_file = os.path.join(dest_dir, filename)
            
            shutil.copy2(source_file, dest_file)
            self.logger.info(f"File transferred: {dest_file}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to transfer file: {e}")
            return False
    
    def transfer_directory(self, source_dir: str, dest_dir: Optional[str] = None) -> bool:
        """
        Transfer directory to iPod device
        
        Args:
            source_dir: Path to source directory
            dest_dir: Destination directory on iPod (optional)
            
        Returns:
            True if successful, False otherwise
        """
        if not self.mount_point or not os.path.exists(self.mount_point):
            self.logger.warning("No device mounted")
            return False
            
        if not os.path.exists(source_dir) or not os.path.isdir(source_dir):
            self.logger.warning(f"Source directory does not exist: {source_dir}")
            return False
            
        try:
            # Determine destination path
            if dest_dir:
                full_dest_dir = os.path.join(self.mount_point, dest_dir)
            else:
                full_dest_dir = os.path.join(self.mount_point, "Music")
            
            os.makedirs(full_dest_dir, exist_ok=True)
            
            for root, dirs, files in os.walk(source_dir):
                rel_path = os.path.relpath(root, source_dir)
                if rel_path == '.':
                    rel_path = ''
                
                dest_subdir = os.path.join(full_dest_dir, rel_path)
                os.makedirs(dest_subdir, exist_ok=True)
                
                for file in files:
                    if file.lower().endswith(('.m4a', '.mp3', '.aac', '.mp4')):
                        src_file = os.path.join(root, file)
                        dst_file = os.path.join(dest_subdir, file)
                        shutil.copy2(src_file, dst_file)
                        self.logger.info(f"Transferred: {dst_file}")
            
            self.logger.info(f"Directory transferred: {source_dir} -> {full_dest_dir}")
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to transfer directory: {e}")
            return False

    def export_track(self, track_info: dict, dest_dir: str) -> Optional[str]:
        """
        Copy a track from iPod to a local directory with metadata-based naming.

        Args:
            track_info: Track dict with keys: file_path, title, artist, album, track_number
            dest_dir: Local destination directory

        Returns:
            Path to the exported file, or None on failure
        """
        src_path = track_info.get("file_path")
        if not src_path or not os.path.exists(src_path):
            self.logger.warning(f"Track file not found on iPod: {track_info.get('title')}")
            return None

        if not os.path.exists(dest_dir):
            os.makedirs(dest_dir, exist_ok=True)

        artist = track_info.get("artist", "Unknown Artist")
        title = track_info.get("title", "Unknown Title")
        album = track_info.get("album", "")
        track_num = track_info.get("track_number", 0)
        ext = os.path.splitext(src_path)[1].lower()

        safe_artist = "".join(c for c in artist if c.isalnum() or c in " ._-'&()").strip() or "Unknown"
        safe_title = "".join(c for c in title if c.isalnum() or c in " ._-'&()").strip() or "Unknown"

        if track_num:
            fname = f"{track_num:02d} - {safe_artist} - {safe_title}{ext}"
        else:
            fname = f"{safe_artist} - {safe_title}{ext}"

        dest_path = os.path.join(dest_dir, fname)

        try:
            shutil.copy2(src_path, dest_path)
            self.logger.info(f"Exported track to: {dest_path}")
            return dest_path
        except Exception as e:
            self.logger.error(f"Failed to export track: {e}")
            return None

    def export_tracks(self, tracks: List[dict], dest_dir: str,
                      progress_callback=None) -> List[Tuple[str, str]]:
        """
        Copy multiple tracks from iPod to a local directory.

        Args:
            tracks: List of track info dicts
            dest_dir: Local destination directory
            progress_callback: Optional callback(progress_int, status_str)

        Returns:
            List of (title, exported_path) tuples for successful exports
        """
        exported = []
        total = len(tracks)

        for i, track in enumerate(tracks):
            if progress_callback:
                progress_callback(int((i / total) * 100), f"Exporting {track.get('title', '?')}...")

            path = self.export_track(track, dest_dir)
            if path:
                exported.append((track.get("title", "?"), path))

        if progress_callback:
            progress_callback(100, f"Exported {len(exported)} of {total} tracks")

        return exported


if __name__ == "__main__":
    ipod = IPodDevice()
    
    # Detect devices
    devices = ipod.detect_devices()
    for device in devices:
        status = "mounted" if device.get("mounted") else "NOT mounted"
        print(f"Found device: {device['name']} ({device['model']}) - {status}")
        if device.get("mount_point"):
            print(f"  Mount point: {device['mount_point']}")
        print(f"  Device: {device['id']}")
    
    # Mount and eject if device found
    if devices:
        device = devices[0]
        if not device.get("mounted"):
            mount_point = ipod.mount_device(device["id"])
            if mount_point:
                print(f"Mounted at: {mount_point}")
                
                info = ipod.get_device_info()
                print(f"Free space: {info.get('free_space', 0) / 1024**2:.1f} MB")
                
                # Unmount
                ipod.unmount_device()
                print("Ejected.")
        else:
            print(f"Already mounted at: {device['mount_point']}")