# Movable Type (r) Open Source (C) 2001-2011 Six Apart, Ltd.
# This program is distributed under the terms of the
# GNU General Public License, version 2.
#
# $Id$

package MT::ObjectDriver::DDL::Pg;

use strict;
use warnings;
use base qw( MT::ObjectDriver::DDL );

sub can_add_column { 1 }
sub can_drop_column { 1 }
sub can_alter_column { 0 }

sub index_defs {
    my $ddl = shift;
    my ($class) = @_;
    my $driver = $class->driver;
    my $dbh = $driver->r_handle;
    my $field_prefix = $class->datasource;
    my $table_name = $class->table_name;
    my $sth = $dbh->prepare(<<SQL)
SELECT cidx.relname as index_name, idx.indisunique, idx.indisprimary, idx.indnatts, idx.indkey, am.amname
FROM pg_index idx
INNER JOIN pg_class cidx ON idx.indexrelid = cidx.oid
INNER JOIN pg_class ctbl ON idx.indrelid = ctbl.oid
INNER JOIN pg_am am ON cidx.relam = am.oid
WHERE ctbl.relname = '$table_name'
SQL
        or return undef;
    $sth->execute or return undef;

    my $defs = {};
    while (my $row = $sth->fetchrow_hashref) {
        next if 1 == $row->{'indisprimary'};

        my $key = $row->{'index_name'};
        next unless $key =~ m/^(mt_)?\Q$field_prefix\E_/;
        $key = 'mt_' . $key unless $key =~ m/^mt_/;

        my $type = $row->{'amname'};
        # ignore fulltext or other unrecognized indexes for now
        next unless $type eq 'btree';

        my $is_unique = $row->{'indisunique'};
        $key =~ s/^mt_\Q$field_prefix\E_//;

        my $indkeys = $row->{'indkey'};
        $indkeys =~ s/\s+/,/g;

        my $sth_att = $dbh->prepare(<<ATTSQL)
SELECT attnum, attname
FROM pg_attribute att
INNER JOIN pg_class ctbl ON att.attrelid = ctbl.oid
WHERE att.attnum IN ($indkeys)
AND ctbl.relname = '$table_name'
ATTSQL
            or next;
        $sth_att->execute or next;
        my $row_att = $sth_att->fetchall_hashref('attnum');
        $sth_att->finish;

        my $cols;
        if ( 1 == $row->{'indnatts'} ) {
            # $indkeys have column's attnum
            my $col = $row_att->{$indkeys}->{'attname'};
            $col =~ s/^\Q$field_prefix\E_//;
            $cols = [ $col ];
        }
        else {
            my @cols;
            for my $indkey ( split ',', $indkeys ) {
                my $col = $row_att->{$indkey}->{'attname'};
                $col =~ s/^\Q$field_prefix\E_//;
                push @cols, $col;
            }
            $cols = \@cols;
        }
        if ( $is_unique ) {
            $defs->{$key} = { 'unique' => 1, 'columns' => $cols };
        }
        else {
            if ((@$cols == 1) && ($key eq $cols->[0])) {
                $defs->{$key} = 1;
            } else {
                $defs->{$key} = { 'columns' => $cols };
            }
        }
    }
    $sth->finish;
    return $defs;
}

sub column_defs {
    my $ddl = shift;
    my ($class) = @_;

    my $table_name = $class->table_name;
    my $field_prefix = $class->datasource;
    my $dbh = $class->driver->r_handle;
    return undef unless $dbh;

    local $dbh->{RaiseError} = 0;
    my $sth = $dbh->column_info( undef, undef, $table_name, undef );
    my $attr = $sth->fetchall_hashref('COLUMN_NAME') or return undef;
    return undef if $sth->err;
    return undef unless %$attr;

    my @pri_keys = $dbh->primary_key( undef, undef, $table_name );
    $attr->{$_}->{PRIMARY_KEY} = 0 for keys %$attr;
    $attr->{$_}->{PRIMARY_KEY} = 1 for @pri_keys;

    my $defs = {};
    foreach my $field_name (keys %$attr) {
        my $col = $attr->{$field_name};
        my $coltype = $ddl->db2type($col->{DATA_TYPE});
        my $colname = lc $col->{COLUMN_NAME};
        $colname =~ s/^\Q$field_prefix\E_//i;
        $defs->{$colname}{type} = $coltype;
        if ( $coltype eq 'string') {
            if (defined $col->{COLUMN_SIZE}) {
                $defs->{$colname}{size} = $col->{COLUMN_SIZE};
            } else {
                $defs->{$colname}{type} = 'text';
            }
        }
        unless ( $col->{NULLABLE} ) {
            $defs->{$colname}{not_null} = 1;
        }
        if ( $col->{PRIMARY_KEY} ) {
            $defs->{$colname}{key} = 1;
        }
    }
    $defs;
}

sub drop_sequence {
    my $ddl = shift;
    my ($class) = @_;
    my $driver = $class->driver;
    my $dbh = $driver->rw_handle;

    # do this, but ignore error since it usually means the
    # sequence didn't exist to begin with
    if (my $col = $class->properties->{primary_key}) {
        ## If it's a complex primary key, use the second half.
        if(ref $col) {
            $col = $col->[1];
        }
        my $def = $class->column_def($col);
        if (exists($def->{auto}) && $def->{auto}) {
        #if ($def->{type} eq 'integer') {
            my $seq = $driver->dbd->sequence_name($class);
            local $dbh->{RaiseError} = 0;
            $dbh->do('DROP SEQUENCE ' . $seq);
        }
    }
    1;
}

sub create_sequence {
    my $ddl = shift;
    my ($class) = @_;

    my $driver = $class->driver;
    my $dbh = $driver->rw_handle;

    if ( my $col = $class->properties->{primary_key} ) {
        ## If it's a complex primary key, use the second half.
        if(ref $col) {
            $col = $col->[1];
        }
        my $def = $class->column_def($col);
        if (exists($def->{auto}) && $def->{auto}) {
        #if ($def->{type} eq 'integer') {
            my $seq = $driver->dbd->sequence_name($class);
            my $table_name = $class->table_name;
            my $field_prefix = $class->datasource;
            my $max_sql = 'SELECT MAX(' . $field_prefix . '_' . $col . ') FROM ' . $table_name;
            my ($start) = $dbh->selectrow_array($max_sql);

            $dbh->do('CREATE SEQUENCE ' . $seq . 
                ($start ? (' START ' . ($start + 1)) : ''));
        }
    }
    1;
}

sub type2db {
    my $ddl = shift;
    my ($def) = @_;
    my $type = $def->{type};
    if ($type eq 'string') {
        return 'varchar(' . $def->{size} . ')';
    } elsif ($type eq 'smallint' ) {
        return 'smallint';
    } elsif ($type eq 'bigint' ) {
        return 'bigint';
    } elsif ($type eq 'boolean') { 
        return 'smallint';
    } elsif ($type eq 'datetime') {
        return 'timestamp';
    } elsif ($type eq 'timestamp') {
        return 'timestamp';
    } elsif ($type eq 'integer') {
        return 'integer';
    } elsif ($type eq 'blob') {
        return 'bytea';
    } elsif ($type eq 'text') {
        return 'text';
    } elsif ($type eq 'float') {
        return 'float';
    }
    Carp::croak("undefined type: " . $type);
}

sub column_sql {
    my $ddl = shift;
    my ($class, $name) = @_;

    # ugly but we need to return the sql to express
    # a column differently based on whether we are declaring
    # the column for creating a table or for altering a column.
    # postgres 7.x does not support the 'not null' and 'default'
    # keywords when altering the column.
    if ((caller(1))[3] =~ m/::create_table_sql$/) {
        my $def = $class->column_def($name);
        return $ddl->SUPER::column_sql($class, $name);
    }

    my $field_prefix = $class->datasource;
    my $def = $class->column_def($name);
    my $type = $ddl->type2db($def);
    return $field_prefix . '_' . $name . ' ' . $type;
}

sub add_column_sql { 
    my $ddl = shift;
    my ($class, $name) = @_;
    my $sql = $ddl->column_sql($class, $name);
    my $driver = $class->driver;
    my $table_name = $class->table_name;
    my $field_prefix = $class->datasource;
    my $dbh = $driver->r_handle;
    my @stmt = ("ALTER TABLE $table_name ADD $sql");

    my $def = $class->column_def($name);
    my $default_value;
    if (exists $def->{default}) {
        $default_value = $def->{default};
        if (($def->{type} =~ m/time/) || $driver->dbd->is_date_col($name)) {
            $default_value = $dbh->quote($driver->dbd->ts2db($default_value));
        } elsif ($def->{type} !~ m/int|float|boolean/) {
            $default_value = $dbh->quote($default_value);
        }
        push @stmt, "ALTER TABLE $table_name ALTER COLUMN ${field_prefix}_${name} SET DEFAULT " . $default_value;
        push @stmt, "UPDATE $table_name SET ${field_prefix}_${name} = $default_value WHERE ${field_prefix}_${name} IS NULL";
    }
    if ($def->{key}) {
        push @stmt, "ALTER TABLE $table_name ADD PRIMARY KEY (${field_prefix}_${name})";
    } elsif (($def->{not_null}) 
          && (70300 < $dbh->{pg_server_version})
          && (exists $def->{default})) {
          #postgresql under 7.3.0 does not support not null in alter table
          #plus,we can't set not null unless there are no rows contains null
        push @stmt, "UPDATE $table_name SET ${field_prefix}_${name} = $default_value WHERE ${field_prefix}_${name} IS NULL";
        push @stmt, "ALTER TABLE $table_name ALTER COLUMN ${field_prefix}_${name} SET NOT NULL";
    }
    return @stmt;
}

sub alter_column_sql {
    my $ddl = shift;
    my ($class, $name) = @_;
    my $sql = $ddl->SUPER::alter_column_sql(@_);
    $sql =~ s/\bMODIFY\b/ALTER COLUMN/;
    return $sql;
}

sub cast_column_sql {
    my $ddl = shift;
    my ($class, $name, $from_def) = @_;
    my $field_prefix = $class->datasource;
    my $def = $class->column_def($name);
    if (($from_def->{type} eq 'text') && ($def->{type} eq 'blob')) {
        return "cast(decode(${field_prefix}_$name, 'escape') as " . $ddl->type2db($def) . ')';
    } elsif (($from_def->{type} eq 'blob') && ($def->{type} eq 'text')) {
        return "cast(encode(${field_prefix}_$name, 'escape') as " . $ddl->type2db($def) . ')';
    } else {
        return "cast(${field_prefix}_$name as " . $ddl->type2db($def) . ')';
    }
}

sub drop_index_sql {
    my $ddl = shift;
    my ($class, $key) = @_;
    my $table_name = $class->table_name;

    my $props = $class->properties;
    my $indexes = $props->{indexes};
    return q() unless exists($indexes->{$key});

    if (ref $indexes->{$key} eq 'HASH') {
        my $idx_info = $indexes->{$key};
        if ($idx_info->{unique} && $ddl->can_add_constraint) {
            return "ALTER TABLE $table_name DROP CONSTRAINT ${table_name}_$key";
        }
    }

    return "DROP INDEX ${table_name}_$key";
}

1;
