#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys
import os.path
import re
import subprocess as subp

if __name__ == "__main__":
  # Make sure we can import stuff
  sys.path.append(os.path.abspath(os.path.dirname(sys.argv[0])+"/../lib-python"))

from prosci.util.pdb import Pdb
from prosci.util.getpdb import getpdb
from prosci.common import NotFoundError, ParsingError


def getchains(pdbcode, chaincodes, dbdir):
    pdbcode = pdbcode.lower()
    a = None
    
    if dbdir is not None:
      try:
        a = getpdb(pdbcode, dbdir)
      except NotFoundError:
        pass
    
    if a is None:
      f=subp.Popen("wget -q -O - http://www.pdb.org/pdb/files/%s.pdb.gz | gunzip 2>/dev/null" % (pdbcode.upper()), shell=True, stdout=subp.PIPE).stdout
      try:
        a = Pdb(pdbcode, f)
      except ParsingError:
        f.close()
        raise NotFoundError("Could not retrieve pdb '%s' from ftp.wwpdb.org"%(pdbcode))
      f.close()
    
    #chains = a.get_chain_codes()
    if not chaincodes:
      return a
    return a.get_chain(chaincodes, include_duplicates=False)


if __name__ == "__main__":
    
    from prosci.shell import Params

    ################################
    # Command line option handling #
    ################################
    
    params = Params(allowed=('help', 'verbose', 'quiet', 'dbdir'), withargument=('dbdir'), helpoption='help', usage="""
          This is a programme to retrieve pdb files from a local database.
          It will create pdb files in the current working directory with
          names in the format 'PDBCODE.pdb'.
          
          USAGE:
            $0 PDBCHAIN ...
      \n""",
      help="""
            PDBCHAIN should be a 5-character id, where the first 4 chars are
            the PDB code and the last char is the chain identifier.

            By default, PDBs are searched for in a directory called "pdb",
            located parallel to the one containing this script. If this
            directory does not exist, PDBs are retrieved from the ftp.wwpdb.org
            server.
            
              OPTIONS:
                --help      Print this message and exit.
                --verbose   Print success messages.
                --quiet     Do not print failure messages.
                --dbdir DIR Search for PDB files in DIR
      \n""")
    
    if len(params.args) < 1:
      params.writeUsage()
      sys.exit(1)
    
    beVerbose = params.isOpt("verbose")
    beQuiet   = params.isOpt("quiet")
    dbdir     = params.getOpt("dbdir")

    if not dbdir:
        dbdir = os.path.abspath(os.path.dirname(sys.argv[0])+"/../pdb")
        if not os.path.isdir(dbdir):
            dbdir=None
    if not beQuiet:
      if dbdir:
        print "Getting PDB file(s) from '%s'..."%(dbdir)
      else:
        print "Getting PDB file(s) from the PDB server..."

    #print "DBDIR", dbdir
    
    def doIt(name, dbdir):
        try:
          mychain = getchains(name[:4], name[4:5], dbdir)
          if not len(mychain):
            sys.stderr.write("ERROR: PDB does not contain specified chain: "+name+"\n")
            return False
          output = open(name+".pdb", "w")
          output.write(str(mychain))
          output.write(str(mychain.ligands))
          output.close()
          return True
        except NotFoundError:
          return False
    
    
    alphanum = re.compile("[a-zA-Z0-9]")
    
    # Reading from STDIN
    #
    if len(params.args) == 1 and params.args[0] == "-":
      status = len(sys.argv)-1
      for arg in sys.stdin:
        #arg = "".join(alphanum.findall(arg))
        if len(arg) == 6:
          arg = arg[:4]+arg[5:]
        if not arg:
          continue
        if len(arg) not in (4, 5):
            sys.stderr.write("ERROR: Argument needs to have length 4 or 5 (counting alphanumeric characters only): "+arg+"\n")
        else:
            if doIt(arg, dbdir):
                status -= 1
                if beVerbose:
                    sys.stderr.write("Getting "+arg+" DONE\n")
            else:
                if not beQuiet:
                    sys.stderr.write("Getting "+arg+" FAILED\n")
      sys.exit(status)
    
    
    # Reading from command line arguments
    #
    status = len(sys.argv)-1
    for arg in params.args:
        #arg = "".join(alphanum.findall(arg))
        if len(arg) == 6:
          arg = arg[:4]+arg[5:]
        if len(arg) not in (4, 5):
            sys.stderr.write("ERROR: Argument needs to have length 4 or 5: "+arg+"\n")
        else:
            if doIt(arg, dbdir):
                status -= 1
                if beVerbose:
                    sys.stderr.write("Getting "+arg+" DONE\n")
            else:
                if not beQuiet:
                    sys.stderr.write("Getting "+arg+" FAILED\n")
    sys.exit(status)
