#!BPY
""" 
Name: 'PostGIS'
Blender: 245
Group: 'Import'
Tooltip: 'Import from PostGIS'
"""


__author__ = "kyk"
__url__ = "http://luxla.net"
__version__ = "1.0"




import Blender, bpy, psycopg2, struct
from Blender import *


WKBPOINT = 1
WKBLINESTRING = 2
WKBPOLYGON = 3
WKBMULTIPOINT = 4
WKBMULTILINESTRING = 5
WKBMULTIPOLYGON = 6
WKBGEOMETRYCOLLECTION = 7



def appendVertexToMesh( mesh, data, size, dim=2 ):
    
    def extend(x, y, z = 0, m = 0):
        mesh.verts.extend(x,y,z)
    
    print "\tStart Vertex"
    
    for ct in range(size):
        extend( *struct.unpack('<' + str(dim) + 'd', data[ct * 8 * dim: ct * 8 * dim + 8 * dim]))
    
    print "\tEnd Vertex"
    
    return data[ct * 8 * dim + 8 *dim:]

def appendEdgeToMesh( mesh, num_edge, start = 0 ):
    print "\tStart Edge"
    mesh.edges.extend( [ (x, x + 1) for x in range(start, start + num_edge) ])
    print "\tEnd Edge"

def WkbPoint( mesh, data, dim = 2 ):
    print 'WKBPoint'
    
    appendVertexToMesh( mesh, data, 1, dim)


def WkbLineString( mesh, data, dim = 2):
    nnumv, = struct.unpack('<1L', data[0:4])
    
    print 'WKBLineString numpoints(%d)' % nnumv
    
    cnumv = len(mesh.verts)
    data = appendVertexToMesh( mesh, data[4:], nnumv, dim)
    appendEdgeToMesh( mesh, nnumv - 1, cnumv)
    
    return data
    

def WkbMultiLineString( mesh, data, dim = 2):
    numl, = struct.unpack('<1L', data[0:4])
    data = data[4:]
    
    print 'WKBMultiLineString numlines(%d)' % numl
    
    for ct in range(numl):
        byteorder, geomtype = struct.unpack('<1c1L',data[0:5])
        if geomtype != WKBLINESTRING:
            raise Error('unknown geometry error')
            
        data = WkbLineString(mesh, data[5:], dim)
    
    return data

def importGeometry( cu, getmesh, dim = 2 ):
    handler = [None, WkbPoint, WkbLineString, None, None, WkbMultiLineString ]
    row = cu.rowcount
    for ct in range(row):
        data = cu.fetchone()
        Window.DrawProgressBar( float(ct) / row, 'Import PostGIS')
        byteorder, geomtype= struct.unpack('<1c1L', data[0][0:5])
        print geomtype
        
        try:
            handler[geomtype]( getmesh(), data[0][5:], dim)
        except IndexError, err:
            print 'unkonown geometry type ' + str(geomtype)
        except Exception, err:
            print err
    
    Window.DrawProgressBar( 1.0 , 'Import PostGIS')


def createSeparateMesh():
    scn = Scene.GetCurrent()
    mesh = bpy.data.meshes.new('PostGIS')
    scn.objects.new(mesh, mesh.name)
    return mesh

class singleMesh:
    def __init__(self):
        scn = Scene.GetCurrent()
        self.mesh = bpy.data.meshes.new('PostGIS')
        scn.objects.new(self.mesh, self.mesh.name)
    
    def __call__(self):
        return self.mesh


registed = Registry.GetKey("PostGISImport",1) or dict()
inikeys = ["host","dbname","user","password","sql", "separate"]
defaultValue = ['localhost', 'postgis', '', '', "select asbinary(the_geom,'NDR') from waterline", 1]
uiDict = dict( (k, Draw.Create(registed.setdefault( k,v))) for k, v in zip(inikeys,defaultValue) )

pgConnection = None


def draw():
    global pgConnection, uiDict
    
    if pgConnection == None:
        uiDict['host'] = Draw.String('Host: ', 0, 20, 200, 200, 30, uiDict['host'].val, 256)
        uiDict['dbname'] = Draw.String('Database: ', 0, 20, 160, 200, 30, uiDict['dbname'].val, 256)
        uiDict['user'] = Draw.String('User: ', 0, 20, 120, 200, 30, uiDict['user'].val, 256)
        uiDict['password'] = Draw.String('Password: ', 0, 20, 80, 200, 30, uiDict['password'].val, 256)
        Draw.PushButton('Connect', 1, 20, 40, 200, 30)
    
    else:
        uiDict['sql'] = Draw.String('Query: ', 0, 20, 120, 300, 30, uiDict['sql'].val, 390)
        uiDict['separate'] = Draw.Toggle('is single?', 0, 20, 80, 200, 30, uiDict['separate'].val)
        Draw.PushButton('Import', 2, 20, 40, 200, 30)

def event(ev):
    global pgConnection, uiDict
    
    if ev == 1:
        
        try:
            pgConnection = psycopg2.connect("dbname=%s host=%s user=%s password=%s" %
                ( uiDict['dbname'], uiDict['host'], uiDict['user'], uiDict['password']
            ))
        except Exception, err:
            print err
            return
        
    elif ev == 2:
        meshfunc = None
        
        if uiDict['separate'].val:
            meshfunc = singleMesh()
        else:
            meshfunc = createSeparateMesh
        
        cursor = pgConnection.cursor()
        
        try:
            cursor.execute(uiDict['sql'].val)
        except Exception, err:
            print err
            return
        
        importGeometry( cursor, meshfunc )
       

Draw.Register(draw, lambda e,v: e == Draw.ESCKEY and not v and Draw.Exit() , event)