require 'dbi'
require 'thread'

# FIXME
module DBI
  module SQL
    def SQL.query?(sql)
      return false unless /^\s*select\b/i =~ sql
      return false if /\s+for\s+update/im =~ sql
      true
    end
  end
end

begin
  require 'postgres' 
  class PGconn
    alias exec async_exec
  end
rescue LoadError
end

class DBIPool
  class DBIPoolError < RuntimeError; end
  class DBIPoolConnNotFound < DBIPoolError; end

  module DBIPoolConn
    private
    def conn
      dbi_pool_info[:conn]
    end

    def add_changed(it)
      dbi_pool_info[:changed].push(it)
    end

    def call_changed(commit)
      if commit
	dbi_pool_info[:changed].each do |it| it.commit end
      else
	dbi_pool_info[:changed].each do |it| it.rollback end
      end
    end

    def dbi_pool_info
      info = Thread.current[:DBIPool]
      raise(DBIPoolConnNotFound) unless info and info[:conn]
      info
    end
  end

  include DBIPoolConn

  @pool = nil
  def self.make_pool(sz, *args)
    @pool = self.new(sz, *args)
  end
  
  def self.pool
    @pool
  end
  
  def self.transaction(&b)
    @pool.transaction(&b)
  end

  def self.commit
    @pool.commit
  end

  def self.rollback
    @pool.rollback
  end

  def initialize(sz, *args)
    @args = args
    @db_name = @args[0]
    @size = sz
    @pool = Queue.new
    sz.times do
      @pool.push(nil)
    end
    @sz_mutex = Mutex.new
  end
  attr_reader :size, :db_name

  def size=(sz)
    @sz_mutex.synchronize do
      sz = 1 if sz <= 0
      diff = sz - @size
      if diff > 0
	diff.times do
	  @pool.push(nil) 
	  @size += 1
	end
      else
	diff.times do
	  conn = @pool.pop
	  conn.disconnect if conn
	  @size -= 1
	end
      end
    end
  end

  def transaction(more_conn=false)
    unless more_conn
      begin
	dbh = conn
	return yield(dbh)
      rescue DBIPoolConnNotFound
      end
    end

    abort = false
    result = nil
    dbh = new_conn
    dbh['AutoCommit'] = false rescue DBI::NotSupportedError
    stack = Thread.current[:DBIPool]
    begin
      Thread.current[:DBIPool] = {:conn => dbh, :changed => []}
      begin
	result = yield(dbh)
	call_changed(true)
      rescue Exception
	abort = true
	call_changed(false)
	dbh.rollback
	raise
      end
    ensure
      dbh.commit unless abort
      @pool.push dbh if dbh
      Thread.current[:DBIPool] = stack
    end
    result
  end

  def new_conn
    loop do
      dbh = @pool.pop || DBI.connect(*@args)
      return dbh if dbh.connected?
      @pool.push nil
    end
  end

  def commit
    conn.commit
  end

  def rollback
    conn.rollback
  end
  alias :abort :rollback

  public :conn
end