#------------------------------------------------------------------------------
#
#  Copyright (c) 2005, Enthought, Inc.
#  All rights reserved.
#
#  This software is provided without warranty under the terms of the BSD
#  license included in enthought/LICENSE.txt and may be redistributed only
#  under the conditions described in the aforementioned license.  The license
#  is also available online at http://www.enthought.com/licenses/BSD.txt
#
#  Thanks for using Enthought open source!
#
#  Author: David C. Morrill
#  Date:   07/01/2005
#
#------------------------------------------------------------------------------

""" Defines the table grid model used by the table editor based on the PyFace
    grid control.
"""

#-------------------------------------------------------------------------------
#  Imports:
#-------------------------------------------------------------------------------

import logging

import wx

import wx.grid as wxg

from traits.api \
    import HasPrivateTraits, Any, Str, Instance, Event, Bool, \
           on_trait_change

from traitsui.api \
    import View, Item, Editor

from traitsui.editors.table_editor \
    import ReversedList

from traitsui.table_filter \
    import TableFilter

from traitsui.ui_traits \
    import SequenceTypes

from pyface.ui.wx.grid.api \
    import GridModel, GridSortEvent

from pyface.ui.wx.grid.trait_grid_cell_adapter \
    import TraitGridCellAdapter

from pyface.timer.api \
    import do_later
import collections


logger = logging.getLogger(__name__)

#-------------------------------------------------------------------------------
#  'TraitGridSelection' class:
#-------------------------------------------------------------------------------

class TraitGridSelection ( HasPrivateTraits ):
    """ Structure for holding specification information.
    """

    # The selected object
    obj = Any

    # The specific trait selected on the object
    name = Str

#-------------------------------------------------------------------------------
#  'TableModel' class:
#-------------------------------------------------------------------------------

class TableModel ( GridModel ):
    """ Model for table data.
    """

    #---------------------------------------------------------------------------
    #  Trait definitions:
    #---------------------------------------------------------------------------

    # The editor that created this model
    editor = Instance( Editor )

    # The current filter
    filter = Instance( TableFilter, allow_none = True )

    # Current filter summary message
    filter_summary = Str( 'All items' )

    # Display the table items in reverse order?
    reverse = Bool( False )

    # Event fired when the table has been sorted
    sorted = Event

    # The current 'auto_add' row
    auto_add_row = Any

    #---------------------------------------------------------------------------
    #  'object' interface:
    #---------------------------------------------------------------------------

    #---------------------------------------------------------------------------
    #  Initializes the object:
    #---------------------------------------------------------------------------

    def __init__ ( self, **traits ):
        """ Initializes the object.
        """
        super( TableModel, self ).__init__( **traits )

        # Attach trait handlers to the list object:
        editor = self.editor
        object = editor.context_object
        name   = ' ' + editor.extended_name

        # Set up listeners for any of the model data changing:
        object.on_trait_change( self._on_data_changed, name, dispatch = 'ui' )
        object.on_trait_change( self.fire_content_changed, name + '.-',
                                dispatch = 'ui' )

        # Set up listeners for any column definitions changing:
        editor.on_trait_change( self.update_columns, 'columns',
                                dispatch = 'ui' )
        editor.on_trait_change( self.update_columns, 'columns_items',
                                dispatch = 'ui' )

        # Initialize the current filter from the editor's default filter:
        self.filter = editor.filter

        # If we are using 'auto_add' mode, create the first 'auto_add' row:
        if editor.auto_add:
            self.auto_add_row = row = editor.create_new_row()
            if row is not None:
                row.on_trait_change( self.on_auto_add_row, dispatch = 'ui' )

    #-- TableModel Interface ---------------------------------------------------

    #---------------------------------------------------------------------------
    #  Disposes of the model when it is no longer needed:
    #---------------------------------------------------------------------------

    def dispose ( self ):
        """ Disposes of the model when it is no longer needed.
        """
        editor = self.editor
        object = editor.context_object
        name   = ' ' + editor.extended_name

        # Remove listeners for any of the model data changing:
        object.on_trait_change( self._on_data_changed, name, remove = True )
        object.on_trait_change( self.fire_content_changed, name + '.-',
                                remove = True )

        # Remove listeners for any column definitions changing:
        editor.on_trait_change( self.update_columns, 'columns', remove = True )
        editor.on_trait_change( self.update_columns, 'columns_items',
                                remove = True )

        # Make sure we have removed listeners from the current filter also:
        if self.filter is not None:
            self.filter.on_trait_change( self._filter_modified, remove = True )

        # Clean-up any links that should be broken:
        self.editor = None

    #---------------------------------------------------------------------------
    #  Returns all model items matching the current filter:
    #---------------------------------------------------------------------------

    def get_filtered_items ( self ):
        """ Returns all model items matching the current filter.
        """
        return self.__filtered_items()

    #---------------------------------------------------------------------------
    #  Returns a single specified item from those items matching the current
    #  filter:
    #---------------------------------------------------------------------------

    def get_filtered_item ( self, index = 0 ):
        """ Returns a single specified item from those items matching the
            current filter.
        """
        try:
            return self.__filtered_items()[ index ]
        except:
            logger.error( 'TableModel error: Request for invalid row %d out of '
                          '%d' % ( index, len( self.__filtered_items() ) ) )
            return None

    #---------------------------------------------------------------------------
    #  Returns the raw, unfiltered index corresponding to a specified filtered
    #  index:
    #---------------------------------------------------------------------------

    def raw_index_of ( self, row ):
        """ Returns the raw, unfiltered index corresponding to a specified
            filtered index.
        """
        if self._filtered_cache is None:
            return row

        return self.editor.filtered_indices[ row ]

    #---------------------------------------------------------------------------
    #  Inserts an object after a specified filtered index:
    #---------------------------------------------------------------------------

    def insert_filtered_item_after ( self, index, item ):
        """ Inserts an object after a specified filtered index.
        """
        mapped_index = 0
        n            = len( self.editor.filtered_indices )
        if index >= n:
            if (index != 0) or (n != 0):
                raise IndexError
        elif index >= 0:
            mapped_index = self.editor.filtered_indices[ index ] + 1
        self.__items().insert( mapped_index, item )
        sorted = self._sort_model()
        if sorted:
            mapped_index = self.__items().index( item )
        self._filtered_cache = None
        return ( mapped_index, sorted )

    #---------------------------------------------------------------------------
    #  Deletes the object at the specified filtered index:
    #---------------------------------------------------------------------------

    def delete_filtered_item_at ( self, index ):
        """ Deletes the object at the specified filtered index.
        """
        if index >= len( self.editor.filtered_indices ):
            raise IndexError

        mapped_index = self.editor.filtered_indices[ index ]
        items        = self.__items()
        object       = items[ mapped_index ]
        del items[ mapped_index ]
        self._filtered_cache = None
        return ( mapped_index, object )

    #---------------------------------------------------------------------------
    #  Updates the table view when columns have been changed:
    #---------------------------------------------------------------------------

    def update_columns ( self ):
        """ Updates the table view when columns have been changed.
        """
        self._columns = None
        self.fire_structure_changed()
        self.editor.refresh()

    #---------------------------------------------------------------------------
    #  Resets any sorting being performed on the underlying model:
    #---------------------------------------------------------------------------

    def no_column_sort ( self ):
        """ Resets any sorting being performed on the underlying model.
        """
        self._sorter = self._filtered_cache = None
        self.column_sorted = GridSortEvent(index = -1)
        #self.fire_structure_changed()

    #-- Event Handlers ---------------------------------------------------------

    #---------------------------------------------------------------------------
    #  Handles the contents of the filter being changed:
    #---------------------------------------------------------------------------

    @on_trait_change( 'filter.+' )
    def _filter_modified ( self ):
        """ Handles the contents of the filter being changed.
        """
        self._filtered_cache = None
        self.fire_structure_changed()
        self.editor.filter_modified()

    #---------------------------------------------------------------------------
    #  Handles the grid firing a 'click' event:
    #---------------------------------------------------------------------------

    def _click_changed ( self, event ):
        """ Handles the grid firing a 'click' event.
        """
        row, col = event

        # Fire the same event on the editor after mapping it to a model object
        # and column name:
        object = self.get_filtered_item( row )
        column = self.__get_column( col )
        self.editor.click = ( object, column )

        # Check to see if the column has a view to display:
        view = column.get_view( object )
        if view is not None:
            column.get_object( object ).edit_traits(
                view   = view,
                parent = self._bounds_for( row, col ) )

        # Invoke the column's click handler:
        column.on_click( object )

    #---------------------------------------------------------------------------
    #  Handles the grid firing a 'dclick' event:
    #---------------------------------------------------------------------------

    def _dclick_changed ( self, event ):
        """ Handles the grid firing a 'dclick' event.
        """
        row, col = event

        # Fire the same event on the editor after mapping it to a model object
        # and column name:
        object = self.get_filtered_item( row )
        column = self.__get_column( col )
        self.editor.dclick = ( object, column )

        # Invoke the column's double-click handler:
        column.on_dclick( object )

    #---------------------------------------------------------------------------
    #  Handles the user modifying the current 'auto_add' mode row:
    #---------------------------------------------------------------------------

    def on_auto_add_row ( self ):
        """ Handles the user modifying the current 'auto_add' mode row.
        """
        object = self.auto_add_row
        object.on_trait_change( self.on_auto_add_row, remove = True )

        self.auto_add_row = row = self.editor.create_new_row()
        if row is not None:
            row.on_trait_change( self.on_auto_add_row, dispatch = 'ui' )

        do_later( self.editor.add_row, object,
                                       len( self.get_filtered_items() ) - 2 )

    #-- GridModel Interface ----------------------------------------------------

    def get_column_count ( self ):
        """ Returns the number of columns for this table.
        """
        return len( self.__get_columns() )

    def get_column_name ( self, index ):
        """ Returns the label of the column specified by the (zero-based) index.
        """
        return self.__get_column( index ).get_label()

    def get_column_size ( self, index ):
        """ Returns the size in pixels of the column indexed by *index*.
            A value of -1 or None means to use the default.
        """
        return self.__get_column( index ).get_width()

    def get_cols_drag_value ( self, cols ):
        """ Returns the value to use when the specified columns are dragged or
            copied and pasted. The parameter *cols* is a list of column indexes.
        """
        return [ self.__get_data_column( col ) for col in cols ]

    def get_cols_selection_value ( self, cols ):
        """ Returns a list of TraitGridSelection objects containing the
            objects corresponding to the grid rows and the traits corresponding
            to the specified columns.
        """
        values = []
        for obj in self.__items( False ):
            values.extend( [ TraitGridSelection(
                                 obj  = obj,
                                 name = self.__get_column_name( col ) )
                             for col in cols ] )
        return values

    def sort_by_column ( self, col, reverse = False ):
        """ Sorts the model data by the column indexed by *col*.
        """
        # Make sure we allow sorts by column:
        factory = self.editor.factory
        if not factory.sortable:
            return

        # Flush the object cache:
        self._filtered_cache = None

        # Cache the sorting information for later:
        self._sorter  = self.__get_column(col).key
        self._reverse = reverse

        # If model sorting is requested, do it now:
        self._sort_model()

        # Indicate the we have been sorted:
        self.sorted = True

        self.column_sorted = GridSortEvent( index    = col,
                                            reversed = reverse )

    def is_column_read_only ( self, index ):
        """ Returns True if the column specified by the zero-based *index* is
            read-only.
        """
        return (not self.__get_column( index ).editable)

    def get_row_count ( self ):
        """ Return the number of rows for this table.
        """
        return len( self.__filtered_items() )

    def get_row_name ( self, index ):
        """ Return the name of the row specified by the (zero-based) *index*.
        """
        return '<undefined>'

    def get_rows_drag_value ( self, rows ):
        """ Returns the value to use when the specified rows are dragged or
            copied and pasted. The parameter *rows* is a list of row indexes.
            If there is only one row listed, then return the corresponding trait
            object. If more than one row is listed, then return a list of objects.
        """
        items = self.__filtered_items()
        return [ items[ row ] for row in rows ]

    def get_rows_selection_value ( self, rows ):
        """ Returns a list of TraitGridSelection objects containing the
            object corresponding to the selected rows.
        """
        items = self.__filtered_items()
        return [ TraitGridSelection( obj = items[ row ] ) for row in rows ]

    def is_row_read_only ( self, index ):
        """ Returns True if the row specified by the zero-based *index* is
            read-only.
        """
        return False

    def get_cell_editor ( self, row, col ):
        """ Returns the editor for the specified cell.
        """

        if self.editor is None:
            return None

        column = self.__get_column( col )
        object = self.get_filtered_item( row )
        editor = column.get_editor( object )
        if editor is None:
            return None

        editor._ui = self.editor.ui

        target, name = column.target_name( object )

        return TraitGridCellAdapter( editor, target, name, '',
                   context = self.editor.ui.context,
                   style   = column.get_style( object ),
                   width   = column.get_edit_width( object ),
                   height  = column.get_edit_height( object ) )

    def get_cell_renderer ( self, row, col ):
        """ Returns the renderer for the specified cell.
        """
        return self.__get_column( col ).get_renderer(
                   self.get_filtered_item( row ) )

    def get_cell_drag_value ( self, row, col ):
        """ Returns the value to use when the specified cell is dragged or
            copied and pasted.
        """
        return self.__get_column( col ).get_drag_value(
                                            self.get_filtered_item( row ) )

    def get_cell_selection_value ( self, row, col ):
        """ Returns a TraitGridSelection object specifying the data stored
            in the table at (*row*, *col*).
        """
        return TraitGridSelection( obj  = self.get_filtered_item( row ),
                                   name = self.__get_column_name( col ) )

    def resolve_selection ( self, selection_list ):
        """ Returns a list of (row, col) grid-cell coordinates that
            correspond to the objects in *selection_list*. For each coordinate,
            if the row is -1, it indicates that the entire column is selected.
            Likewise coordinates with a column of -1 indicate an entire row
            that is selected. For the TableModel, the objects in
            *selection_list* must be TraitGridSelection objects.
        """
        items = self.__filtered_items()
        cells = []
        for selection in selection_list:
            row = -1
            if selection.obj is not None:
                try:
                    row = items.index( selection.obj )
                except ValueError:
                    continue

            column = -1
            if selection.name != '':
                column = self._get_column_index_by_trait( selection.name )
                if column is None:
                    continue

            cells.append( ( row, column ) )

        return cells

    def get_cell_context_menu ( self, row, col ):
        """ Returns a Menu object that generates the appropriate context
            menu for this cell.
        """
        column = self.__get_column( col )
        menu   = column.get_menu( self.get_filtered_item( row ) )
        editor = self.editor

        if menu is None:
            menu = editor.factory.menu

        if menu is None:
            menu_name = editor.factory.menu_name
            if menu_name :
                menu = getattr(self.editor.object, menu_name, None)

        if menu is not None:
            editor.prepare_menu( row, column )

            return ( menu, editor )

        return None

    def get_value ( self, row, col ):
        """ Returns the value stored in the table at (*row*, *col*).
        """
        object = self.get_filtered_item( row )
        if object is self.auto_add_row:
            return ''

        value   = self.__get_column( col ).get_value( object )
        formats = self.__get_column_formats( col )

        if (value is not None) and (formats is not None):
            format = formats.get( type( value ) )
            if format is not None:
                try:
                    if isinstance( format, collections.Callable):
                        value = format( value )
                    else:
                        value = format % value
                except:
                    pass

        return value

    def is_valid_cell_value ( self, row, col, value ):
        """ Tests whether *value* is valid for the cell at (*row*, *col*).
        Returns True if value is acceptable, and False otherwise. """
        return self.__get_column( col ).is_droppable(
                   self.get_filtered_item( row ), value )

    def is_cell_empty ( self, row, col ):
        """ Returns True if the cell at (*row*, *col*) has a None value, and
            False otherwise.
        """
        return (self.get_value( row, col ) is None)

    def is_cell_read_only ( self, row, col ):
        """ Returns True if the cell at (*row*, *col*) is read-only, and False
            otherwise.
        """
        return (not self.__get_column( col ).is_editable(
                   self.get_filtered_item( row ) ))

    def get_cell_bg_color ( self, row, col ):
        """ Returns a wxColour object specifying the background color
            of the specified cell.
        """
        return self.__get_column( col ).get_cell_color(
                   self.get_filtered_item( row ) )

    def get_cell_text_color ( self, row, col ):
        """ Returns a wxColour object specifying the text color of the
            specified cell.
        """
        column = self.__get_column(col)
        item = self.get_filtered_item(row)
        return column.get_text_color(item)

    def get_cell_font ( self, row, col ):
        """ Returns a wxFont object specifying the font of the specified cell.
        """
        return self.__get_column( col ).get_text_font(
                   self.get_filtered_item( row ) )

    def get_cell_halignment ( self, row, col ):
        """ Returns a string specifying the horizontal alignment of the
            specified cell.

            Returns 'left' for left alignment, 'right' for right alignment,
            or 'center' for center alignment.
        """
        return self.__get_column( col ).get_horizontal_alignment(
                   self.get_filtered_item( row ) )

    def get_cell_valignment ( self, row, col ):
        """ Returns a string specifying the vertical alignment of the
            specified cell.

            Returns 'top' for top alignment, 'bottom' for bottom alignment,
            or 'center' for center alignment.
        """
        return self.__get_column( col ).get_vertical_alignment(
                   self.get_filtered_item( row ) )

    #---------------------------------------------------------------------------
    #  Protected 'GridModel' interface:
    #---------------------------------------------------------------------------

    def _insert_rows ( self, pos, num_rows ):
        """ Inserts *num_rows* at *pos*; fires an event only if a factory
        method for new rows is defined or the model is not empty. Otherwise,
        it returns 0.
        """
        count = 0

        factory = self.editor.factory.row_factory
        if factory is None:
            items = self.__items( False )
            if len( items ) > 0:
                factory = items[0].__class__

        if factory is not None:
            new_data = [ x for x in [ factory() for i in range( num_rows ) ]
                         if x is not None ]
            if len( new_data ) > 0:
                count           = self._insert_rows_into_model( pos, new_data )
                self.rows_added = ( 'added', pos, new_data )

        return count

    def _delete_rows ( self, pos, num_rows ):
        """ Removes rows *pos* through *pos* + *num_rows* from the model.
        """
        row_count = self.get_rows_count()
        if (pos + num_rows) > row_count:
            num_rows = row_count - pos

        return self._delete_rows_from_model( pos, num_rows )

    def _set_value ( self, row, col, value ):
        """ Sets the value of the cell at (*row*, *col*) to *value*.

            Raises a ValueError if the value is vetoed or the cell at
            the specified position does not exist.
        """
        new_rows = 0
        column   = self.__get_column( col )
        obj      = None
        try:
            obj = self.get_filtered_item( row )
        except:
            # Add a new row:
            new_rows = self._insert_rows( self.get_row_count(), 1 )
            if new_rows > 0:
                # Now set the value on the new object:
                try:
                    obj = self.get_filtered_item( self.get_row_count() - 1 )
                except:
                    # fixme: what do we do in this case? veto the set somehow?
                    # raise an exception?
                    pass

        if obj is not None:
            self._set_data_on_row( obj, column, value )

        return new_rows

    def _move_column ( self, frm, to ):
        """ Moves a specified **frm** column to before the specified **to**
            column. Returns **True** if successful; **False** otherwise.
        """
        to_column = None
        if to < len( self.__get_columns() ):
            to_column = self.__get_column( to )

        return self.editor.move_column( self.__get_column( frm ), to_column )

    #---------------------------------------------------------------------------
    #  Protected interface:
    #---------------------------------------------------------------------------

    def _set_data_on_row ( self, row, column, value ):
        """ Sets the cell specified by (*row*, *col*) to *value, which
            can be either a member of the row object, or a no-argument method
            on that object.
        """
        column.set_value( row, value )

    def _insert_rows_into_model ( self, pos, new_data ):
        """ Inserts the given new rows into the model.
        """
        raw_pos = self.raw_index_of( pos )
        self.__items()[ raw_pos: raw_pos ] = new_data

    def _delete_rows_from_model ( self, pos, num_rows ):
        """ Deletes the specified rows from the model.
        """
        raw_rows = [ self.raw_index_of( i )
                     for i in range( pos, pos + num_rows ) ]
        raw_rows.sort()
        raw_rows.reverse()
        items = self.__items()
        for row in raw_rows:
            del items[ row ]

        return num_rows

    #---------------------------------------------------------------------------
    #  Trait event handlers:
    #---------------------------------------------------------------------------

    def _on_data_changed ( self ):
        """ Forces the grid to refresh when the underlying list changes.
        """
        # Invalidate the current cache (if any):
        self._filtered_cache = None

        self.fire_structure_changed()

    def _mouse_cell_changed ( self, new ):
        """ Handles the user mousing over a specified cell.
        """
        row, col = new
        column   = self.__get_column( col )
        object   = self.get_filtered_item( row )

        # Update the tooltip if necessary:
        tooltip = column.get_tooltip( object )
        if tooltip != self._tooltip:
            self._tooltip = tooltip
            self.editor.grid._grid_window.SetToolTip( wx.ToolTip( tooltip ) )

        if column.is_auto_editable( object ):
            x, y, dx, dy = self._bounds_for( row, col )
            if column.is_editable( object ):
                view = View( Item( name       = column.name,
                                   editor     = column.get_editor( object ),
                                   style      = column.get_style(  object ),
                                   show_label = False,
                                   padding    = -4 ),
                             kind   = 'info',
                             width  = dx,
                             height = dy )
            else:
                view = column.get_view( object )
                if view is None:
                    return

            column.get_object( object ).edit_traits(
                view   = view,
                parent = ( x, y, dx, dy ) )

    #---------------------------------------------------------------------------
    #  Private interface:
    #---------------------------------------------------------------------------

    def _bounds_for ( self, row, col ):
        """ Returns the coordinates and size of the specified cell in the form:
            ( x, y, dx, dy ).
        """
        grid         = self.editor.grid
        coords       = wxg.GridCellCoords( row, col )
        x, y, dx, dy = grid._grid.BlockToDeviceRect( coords, coords )
        x, y         = grid._grid_window.ClientToScreenXY( x, y )

        return ( x, y, dx, dy )

    def _sort_model ( self ):
        """ Sorts the underlying model if that is what the user requested.
        """
        editor = self.editor
        sorted = (editor.factory.sort_model and (self._sorter is not None))
        if sorted:
            items = self.__items( False )[:]
            items.sort(key=self._sorter)
            if self.reverse ^ self._reverse:
                items.reverse()
            editor.value = items
        return sorted

    def __items ( self, ordered = True ):
        """ Returns the raw list of model objects.
        """
        result = self.editor.value
        if not isinstance( result, SequenceTypes ):
            return [ result ]

        if ordered and self.reverse:
            return ReversedList( result )

        return result

    def __filtered_items ( self ):
        """ Returns the list of all model objects that pass the current filter.
        """
        fc = self._filtered_cache
        if fc is None:
            items  = self.__items()
            filter = self.filter
            if filter is None:
                nitems = [ nitem for nitem in enumerate( items ) ]
                self.filter_summary = 'All %s items' % len( nitems )
            else:
                if not isinstance( filter, collections.Callable):
                    filter = filter.filter
                nitems = [ nitem for nitem in enumerate( items )
                           if list(filter( nitem[1] )) ]
                self.filter_summary = '%s of %s items' % ( len( nitems ),
                                                           len( items ) )
            sorter = self._sorter
            if sorter is not None:
                nitems.sort(key=lambda x: sorter(x[1]))
                if self._reverse:
                    nitems.reverse()

            self.editor.filtered_indices = [ x[0] for x in nitems ]
            self._filtered_cache = fc    = [ x[1] for x in nitems ]
            if self.auto_add_row is not None:
                self._filtered_cache.append( self.auto_add_row )

        return fc

    def __get_data_column ( self, col ):
        """ Returns a list of model data from the column indexed by *col*.
        """
        column = self.__get_column( col )
        return [ column.get_value( item ) for item in self.__filtered_items() ]

    def __get_columns ( self ):
        columns = self._columns
        if columns is None:
            self._columns = columns = [ c for c in self.editor.columns
                                                if c.visible ]
        return columns

    def __get_column ( self, col ):
        try:
            return self.__get_columns()[ col ]
        except:
            return self.__get_columns()[0]

    def __get_column_name ( self, col ):
        return self.__get_column( col ).name

    def __get_column_formats ( self, col ):
        return None   # Not used/implemented currently

    def _get_column_index_by_trait ( self, name ):
        for i, col in enumerate( self.__get_columns() ):
            if name == col.name:
                return i
