import _parted
import sys
import types

SECTOR_SIZE = 512

###############################################################################
# PedSequence

class PedSequence:
    """This class implements the *_get_next functions by creating
    a special sequence type. This allows sequences of parted
    objects to be treated just like any other Python immutable
    sequence. The sequence avoids actually instantiating any
    parted objects until they are needed, however, and will
    handle growing and shrinking as parted objects are created or
    destroyed properly (for example: a sequence referring to all
    partitions on a disk will remain correct even as partitions
    are removed or added).
    """

    def __init__(self, start_item_func, *start_item_params):
        self.start_item_func = start_item_func
        if len(start_item_params) > 0:
            self.start_item_param = start_item_params[0]
        else:
            self.start_item_param = None
        self.reset()

    def reset(self):
        self.list = []
        firstitem = self.start_item_func(self.start_item_param)
        if firstitem is not None:
            self.list.append(firstitem)

    def enumerate(self):
        if len(self.list) > 0:
            next_item = self.list[-1]._get_next()
            while next_item:
                self.list.append(next_item)
                next_item = next_item._get_next()

    def __len__(self):
        self.reset()
        self.enumerate()
        return len(self.list)

    def __getitem__(self, key):
        if not isinstance(key, types.IntType):
            raise IndexError, "sequence index not an integer"

        self.reset()

        if key < 0:
            self.enumerate()
        else:
            if len(self.list) < (key + 1) and len(self.list) > 0:
                next_item = self.list[-1]._get_next()
                while next_item and len(self.list) < (key + 1):
                    self.list.append(next_item)
                    next_item = next_item._get_next()

            if len(self.list) < (key + 1):
                raise IndexError, "sequence index out of range"

        return self.list[key]

    def __getslice__(self, i, j):
        self.reset()
        if i < 0 or j == sys.maxint or j < 0:
            self.enumerate()
        else:
            # Force enumeration to the last index.
            dummy = self[j]

        return self.list[i:j]

###############################################################################
# PEDDEVICE

# Enum DeviceType
DEVICE_UNKNOWN	= 0
DEVICE_SCSI	= 1
DEVICE_IDE	= 2
DEVICE_DAC960	= 3

class Device:
    """Represents a storage device.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def _get_next(self):
        return _device_get_next(self)

    def get_model(self):
        """Get a description of the hardware manufacturer and model.
        """

        return _parted.device_get_model(self._o)

    def get_path(self):
        """Get the block device, e.g. /dev/sdb.
        """

        return _parted.device_get_path(self._o)

    def get_type(self):
        return _parted.device_get_type(self._o)

    def get_sector_size(self):
        return _parted.device_get_sector_size(self._o)

    def get_heads(self):
        return _parted.device_get_heads(self._o)

    def get_sectors(self):
        return _parted.device_get_sectors(self._o)

    def get_cylinders(self):
        return _parted.device_get_cylinders(self._o)

    def get_geom_known(self):
        return _parted.device_get_geom_known(self._o)

    def get_host(self):
        return _parted.device_get_host(self._o)

    def get_did(self):
        return _parted.device_get_did(self._o)

    def get_length(self):
        """Get the size of the device in sectors.
        """

        return _parted.device_get_length(self._o)

    def get_open_count(self):
        return _parted.device_get_open_count(self._o)

    def get_dirty(self):
        return _parted.device_get_dirty(self._o)

    # Action Methods

    def open(self):
        """Attempts to open dev, to allow use uf read(), write()
	and sync(). Returns zero on failure.
        """

        return _parted.device_open(self._o)

    def close(self):
        """Closes the device.
        """

        return _parted.device_close(self._o)

    def read(self, start, count):
        """Reads count sectors from the device, beginning with
        sector start.
        """

        return _parted.device_read(self._o, start, count)

    def write(self, buffer, start, count):
        """Writes count sectors to the device, beginning with
        sector start.
        """

        return _parted.device_write(self._o, buffer, start, count)

    def sync(self):
        """Flushes the write cache.
        """

        return _parted.device_sync(self._o)

    def disk_open(self):
        """Constructs a PedDisk object from dev, and reads the
	partition table. Returns None on failure. WARNING: this
	can modify dev->cylinders, dev->heads and dev->sectors,
	because the partition table might indicate that the
	existing values were incorrect.
        """

        r = _parted.disk_open(self._o)
        if r is None:
            return None
        else:
            return Disk(r)

    def disk_create(self, type):
        """Creates a partition table on this device and returns a
        Disk object for it.
        """

        return Disk(_parted.disk_create(self._o, type._o))

def _device_get_next(dev=None):
    """Returns the next device that was detected by device_probe_all(), or
    calls to ped_device_get_next().  If dev is None, returns the first
    device.  Returns None if dev is the last device.
    """

    if dev:
        r = _parted.device_get_next(dev._o)
    else:
        r =_parted.device_get_next(None)

    if r:
        return Device(r)
    else:
        return None

###############################################################################
# PEDDISK, PEDDISKTYPE

class DiskType:
    def __init__(self, obj):
        self._o = obj

    def _get_next(self):
        return _disk_type_get_next(self)

    def get_name(self):
        """Get the name of the type.
        """

        return _parted.disk_type_get_name(self._o)

    def get_ops(self):
        return _parted.disk_type_get_ops(self._o)

def _disk_type_get_next(type=None):
    if type:
        r = _parted.disk_type_get_next(type._o)
    else:
        r = _parted.disk_type_get_next(None)

    if r:
        return DiskType(r)
    else:
        return None

class Disk:
    """Represents a device + partition table.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def get_dev(self):
        return Device(_parted.disk_get_dev(self._o))

    def get_type(self):
        """The type of the disk label.
        """

        return DiskType(_parted.disk_get_type(self._o))

    def get_part_list(self):
        """The list of partitions on the disk.
        """

        return PedSequence(_disk_next_partition, self)

    # Action Methods

    def close(self):
        """Closes the disk.
        """

        return _parted.disk_close(self._o)

    def read(self):
        """Reads the partition table from the disk.
        """

        return _parted.disk_read(self._o)

    def write(self):
        """Writes the partition table to the disk.
        """

        return _parted.disk_write(self._o)

    def add_partition(self, part):
	"""Adds "part" to "disk". "part"'s geometry may be
	changed in this process. Sorry, I know this is bad, but
	there is no other way. Some brain-dead partition table
	systems, not mentionioning any Micros^H^H^H^H^H^Hnames
	require certain alignments which can't be pre-calculated.
	"part" is assigned a number in this process.
        """

        part.disk = self
        return _parted.disk_add_partition(self._o, part._o)

    def delete_partition(self, part):
        """Removes part from the disk, destroying part.
        """

        return _parted.disk_delete_partition(self._o, part._o)

    def delete_all(self):
        """Removes an destroys all partitions on the disk.
        """

        return _parted.disk_delete_all(self._o)

    def set_partition_geom(self, part, start, end):
        """Sets the geometry of part.
        """

        return _parted.disk_set_partition_geom(self._o, part._o, start, end)

    def maximize_partition(self, part):
        """Grows part's geometry to the maximum possible.
        """

        return _parted.disk_maximize_partition(self._o, part._o)

    def get_max_partition_geometry(self, part):
        """Returns the maximum geometry part can be grown to.
        """

        return Geometry(_parted.disk_get_max_paritition_geometry(self._o,
                                                                 part._o))

    def minimize_extended_partition(self):
        """Reduces the extended partition on the disk to the
        minimum possible."""

        return _parted.disk_minimize_extended_partition(self._o)

    def next_partition(self, part):
        return _disk_next_partition(self, part)

    def get_partition(self, num):
        """Returns the partition numbered num.
        """

        r = _parted.disk_get_partition(self._o, num)
        if r:
            return Partition(self, None, None, None, None, r)
        else:
            return None

    def get_partition_by_sector(self, sect):
        """Returns the partition that owns sect.
        """

        return Partition(self, None, None, None, None,
                         _parted.disk_get_partition_by_sector(self._o, sect))

    def get_extended_partition(self):
        """Returns the extended partition.
        """

        return Partition(self, None, None, None, None,
                         _parted.disk_get_extended_partition(self._o))

def _disk_next_partition(disk, part=None):
    if part is not None:
        r = _parted.disk_next_partition(disk._o, part._o)
    else:
        r = _parted.disk_next_partition(disk._o)

    if r:
        return Partition(disk, None, None, None, None, r)
    else:
        return None

###############################################################################
# PEDGEOMETRY

class Geometry:
    """Represents a continuous region on a device.
    """

    def __init__(self, obj):
        self._o = obj

    # Accessor Methods

    def get_disk(self):
        """Returns the disk.
        """

        return Disk(_parted.geometry_get_disk(self._o))

    def get_start(self):
        """Returns the start of the region in sectors.
        """

        return _parted.geometry_get_start(self._o)

    def get_length(self):
        """Returns the length of the region, in sectors.
        """

        return _parted.geometry_get_length(self._o)

    def get_end(self):
        """Returns the end of the region in sectors.
        """

        return _parted.geometry_get_end(self._o)

    # I'm not going to bind anymore of this one just yet, as it
    # appears no one is using any of the action methods.

###############################################################################
# PEDPARTITION, PEDPARTITIONTYPE

# Enum PartitionType:
PED_PARTITION_PRIMARY		= 0x00
PED_PARTITION_LOGICAL		= 0x01
PED_PARTITION_EXTENDED		= 0x02
PED_PARTITION_FREESPACE		= 0x04
PED_PARTITION_METADATA		= 0x08

class Partition:
    """Represents a partition. This is basically a Geomtry plus
    some attributes.
    """

    def __init__(self, disk, type, fs_type, start, end,
                 obj=None):
        """Creates a new Partition on disk.
        """

        # Remember disk for iteration later.

        if disk is None:
            raise ValueError, "no disk object passed in constructor"
        self.disk = disk

        if obj:
            self._o = obj
        else:
            if type == PED_PARTITION_EXTENDED:
                r = _parted.partition_new(disk._o, type, None,
                                          start, end)
            else:
                r = _parted.partition_new(disk._o, type, fs_type._o,
                                          start, end)
            if r:
                self._o = r
            else:
                raise RuntimeError, "unable to create partition"

    # Accessor methods

    def _get_next(self):
        """Retrieve next partition on the same disk.
        """

        if self.disk is None:
            return None
        else:
            return _disk_next_partition(self.disk, self)

    def get_geom(self):
        """Geometry of the partition.
        """

        return Geometry(_parted.partition_get_geom(self._o))

    def get_num(self):
        """The partition number.
        """

        return _parted.partition_get_num(self._o)

    def get_hidden(self):
        """Whether the partition is hidden.
        """

        return _parted.partition_get_hidden(self._o)

    def get_type(self):
        """The partition type.
        """

        return _parted.partition_get_type(self._o)

    def get_part_list(self):
        """List of logical partitions in an extended partition.
        """

        return Partition(self.disk, None, None, None, None,
                         _parted.partition_get_part_list(self._o))

    def get_bootable(self):
        """Whether the partition is the one that should be booted.
        """

        return _parted.partition_get_bootable(self._o)

    def set_bootable(self, value):
        """Set the bootable flag to VALUE.
        """

        return _parted.partition_set_bootable(self._o, value)
        
    def get_fs_type(self):
        """Type of filesystem on the partition.
        """

        ped_object = _parted.partition_get_fs_type(self._o)
        if ped_object is None:
            return None
        else:
            return FileSystemType(ped_object)

    def set_fs_type(self, type):
        """Set the file system type on the partition to TYPE.
        """

        return _parted.partition_set_fs_type(self._o, type._o)
    
    def is_active(self):
        """Whether the partition is active.
        """

        return _parted.partition_is_active(self._o)

    def is_busy(self):
        """Whether the partition is busy (i.e. mounted).
        """

        return _parted.partition_is_busy(self._o)

    # Member Set Methods

    def set_system(self, fs_type):
        """Sets the system type.
        """

        return _parted.partition_set_system(self._o, fs_type._o)

    # Action Methods

    def destroy(self):
        """Destroys a partition. Should not be called on a a
        partition that is in the partition table. Use
        Disk.delete_partition() instead.
        """

        return _parted_partition_destroy(self._o)

###############################################################################
# PEDFILESYSTEM, PEDFILESYSTEMTYPE

class FileSystemType:
    def __init__(self, obj):
        self._o = obj

    def _get_next(self):
        return _file_system_type_get_next(self)

    def get_name(self):
        """Name of the filesystem type.
        """

        return _parted.file_system_type_get_name(self._o)

    def get_ops(self):
        return FileSystemType(_parted.file_system_type_get_ops(self._o))

def _file_system_type_get_next(fs_type=None):
    if fs_type:
        r = _parted.file_system_type_get_next(fs_type._o)
    else:
        r = _parted.file_system_type_get_next(None)

    if r:
        return FileSystemType(r)
    else:
        return None

class FileSystem:
    """Represents a filesystem, associated with a Geometry, NOT a
    Partition.
    """

    def __init__(self, geom, type, obj=None):
        if obj:
            self._o = obj
        else:
            self._o = _parted.file_system_create(geom._o, type._o)

    # Accessor Methods

    def get_type(self):
        return FileSystemType(_parted.file_system_get_type(self._o))

    def get_geom(self):
        """Where the filesystem actually is.
        """

        return Geometry(_parted.file_system_get_geom(self._o))

    # Action Methods

    def close(self):
        """Closes the filesystem.
        """

        return _parted.file_system_close(self._o)

    def check(self):
        """Checks filesystem for errors.
        """

        return _parted.file_system_check(self._o)

    def copy(self, geom):
        """Creates a new filesystem of the same type on geom, and
        copies the contents of this fs into the new one.
        """

        return _parted.file_system_copy(self._o, geom._o)

    def resize(self, geom):
        """Resize the filesystem to the new geom.
        """

        return _parted.file_system_resize(self._o, geom._o)

def file_system_probe(geom):
    """Attempts to detect a filesystem on geom and returns it.
    """

    return _parted.file_system_probe(geom._o)

def file_system_open(geom):
    """Opens a filesystem on geom.
    """

    return _parted.file_system_open(geom._o)

###############################################################################
# Exceptions

# Enum PedExceptionType
EXCEPTION_INFORMATION=1
EXCEPTION_WARNING=2
EXCEPTION_ERROR=3
EXCEPTION_FATAL=4
EXCEPTION_BUG=5
EXCEPTION_NO_FEATURE=6

# Enum PedExceptionOption
EXCEPTION_UNHANDLED=0
EXCEPTION_FIX=1
EXCEPTION_YES=2
EXCEPTION_NO=4
EXCEPTION_OK=8
EXCEPTION_RETRY=16
EXCEPTION_IGNORE=32
EXCEPTION_CANCEL=64

# These are used by the default exception handler.
type_strings = {
    EXCEPTION_INFORMATION: "information",
    EXCEPTION_WARNING:     "warning",
    EXCEPTION_ERROR:       "error",
    EXCEPTION_FATAL:       "fatal error",
    EXCEPTION_BUG:         "parted bug",
    EXCEPTION_NO_FEATURE:  "feature not implemented"
    }

response_strings = {
    EXCEPTION_FIX:    ["Fix", "f"],
    EXCEPTION_YES:    ["Yes", "y"],
    EXCEPTION_NO:     ["No", "n"],
    EXCEPTION_OK:     ["OK", "o"],
    EXCEPTION_RETRY:  ["Retry", "r"],
    EXCEPTION_IGNORE: ["Ignore", "i"],
    EXCEPTION_CANCEL: ["Cancel", "c"]
    }

class PedException:
    """This class encapsulates parted exceptions.  Exceptions in
    parted are not, strictly speaking, exceptions; parted functions
    will either invoke an error handler callback or return an error
    code, depending on the preference of the application.  In Python,
    when a parted exception is thrown, an instance of this class will
    be created to encapsulate the information in the exception.  This
    is used to provide an easy-to-use interface for applications to
    use when registering exception handlers.
    """

    def __init__(self, message, type, options):
        self.type = type
        self.options = options
        self.message = message
        self.iscaught = 0

        # Trigger exception if options are invalid.
        dummy = self.get_options()

    def get_options(self):
        """Return a tuple of possible options to return.
        """

        optionlist = []
        for pos_option in (EXCEPTION_FIX, EXCEPTION_YES, EXCEPTION_NO,
                           EXCEPTION_OK, EXCEPTION_RETRY,
                           EXCEPTION_IGNORE, EXCEPTION_CANCEL):
            if self.options & pos_option:
                optionlist.append(pos_option)
        if length(optionlist) < 1:
            raise AttributeError, "no valid option set"
        return optionlist

    def get_type(self):
        return self.type

    def get_message(self):
        return self.message

def _lowlevel_exception_handler(message, type, options):
    """Hook our exception handling interface into the low-level parted
    exception system."""

    if pedobject is None:
        return EXCEPTION_UNHANDLED

    try:
        ped_exception = PedException(message, type, options)
    except AttributeError:
        return EXCEPTION_UNHANDLED

    return call_handler(ped_exception)

def _default_handler(ped_exception):
    """Default parted exception handler.  It writes the error and a
    list of options to standard error, and expects the letter of one
    response to be entered on standard input.
    """
    
    sys.stderr.write("parted system: %s: %s\n" %
                     (type_strings[ped_exception.get_type()],
                      ped.exception.get_message()))

    options = ped_exception.get_options()
    for option in options:
        optionstr = response_strings[option]
        sys.stderr.write("%s (%s)", optionstr[0], optionstr[1])
    sys.stderr.write("? ")

    response_value = 0
    while not response_value:
        response = f.readline()[0]
        for option in options:
            optionstr = response_strings[option]
            if response == optionstr[1]:
                response_value = option

    return response_value

################################################################################
# Top-level functions and globals

_exception_handler = _default_handler

def init():
    return _parted.init()

def done():
    """Shut down the underlying parted subsystem.
    """

    return _parted.done()

def device_probe_all():
    """Attempts to detect all devices.
    """

    return _parted.device_probe_all()

def device_get(name):
    """Gets the device "name", where name is usually the block device (eg
    /dev/sdb).  If the device wasn't detected with device_probe_all(),
    an attempt will be made to detect it again.
    """

    return Device(_parted.device_get(name))

def disk_type_get(name):
    return DiskType(_parted.disk_type_get(name))

def file_system_type_get(name):
    """Returns the FileSystemType corresponding to name.
    """

    return FileSystemType(_parted.file_system_type_get(name))

def call_handler(ped_exception):
    return (_exception_handler)(ped_exception)

def register_handler(handler_func):
    _exception_handler = handler_func

def unregister_handler():
    _exception_handler = _default_handler

def get_devices():
    return PedSequence(_device_get_next)

def get_disk_types():
    return PedSequence(_disk_type_get_next)

def get_file_system_types():
    return PedSequence(_file_system_type_get_next)