#! /usr/bin/python

from sys import argv, stdout, stderr, exit
import errno as E
import fcntl as F
import os as OS
import subprocess as S
import select as SEL
import signal as SIG

import catacomb as C
import sqlite3 as SQL

class ExpectedError (Exception):
  pass

CONNINIT_SQL = """
PRAGMA foreign_keys = on;
"""

SETUP_SQL = """
PRAGMA journal_mode = wal;

CREATE TABLE top
        (kind TEXT NOT NULL,            -- `gf2x'
         groupdesc TEXT NOT NULL,
         g TEXT NOT NULL,
         x TEXT NOT NULL,
         m TEXT NOT NULL,               -- g^m = 1
         n TEXT DEFAULT NULL);          -- g^n = x

CREATE TABLE progress
        (p TEXT PRIMARY KEY NOT NULL,   -- p|m, p prime
         e INT NOT NULL,                -- e = v_p(m)
         k INT NOT NULL DEFAULT(0),     -- 0 <= k <= e
         n TEXT NOT NULL DEFAULT(0),    -- (g^{m/p^k})^n = x^{m/p^k}
         dpbits INT NOT NULL);          -- 0 for sequential
CREATE UNIQUE INDEX progress_by_p_k ON progress (p, k);

CREATE TABLE workers
        (pid INT PRIMARY KEY NOT NULL,
         p TEXT NOT NULL,
         k INT NOT NULL,
         FOREIGN KEY (p, k) REFERENCES progress (p, k));
CREATE INDEX workers_by_p ON workers (p, k);

CREATE TABLE points
        (p TEXT NOT NULL,
         k INT NOT NULL,
         z TEXT NOT NULL,               -- g^a x^b = z
         a TEXT NOT NULL,
         b TEXT NOT NULL,
         PRIMARY KEY (p, k, z),
         FOREIGN KEY (p, k) REFERENCES progress (p, k));
"""

GROUPMAP = {}

class GroupClass (type):
  def __new__(cls, name, supers, dict):
    ty = super(GroupClass, cls).__new__(cls, name, supers, dict)
    try: name = ty.NAME
    except AttributeError: pass
    else: GROUPMAP[name] = ty
    return ty

class BaseGroup (object):
  __metaclass__ = GroupClass
  def __init__(me, desc):
    me.desc = desc
  def div(me, x, y):
    return me.mul(x, me.inv(y))

class BinaryFieldUnitGroup (BaseGroup):
  NAME = 'gf2x'
  def __init__(me, desc):
    super(BinaryFieldUnitGroup, me).__init__(desc)
    p = C.GF(desc)
    if not p.irreduciblep(): raise ExpectedError, 'not irreducible'
    me._k = C.BinPolyField(p)
    me.order = me._k.q - 1
  def elt(me, x):
    return me._k(C.GF(x))
  def pow(me, x, n):
    return x**n
  def mul(me, x, y):
    return x*y
  def inv(me, x):
    return x.inv()
  def idp(me, x):
    return x == me._k.one
  def eq(me, x, y):
    return x == y
  def str(me, x):
    return str(x)

def getgroup(kind, desc): return GROUPMAP[kind](desc)

def factor(n):
  ff = []
  proc = S.Popen(['./factor', str(n)], stdout = S.PIPE)
  for line in proc.stdout:
    pstr, estr = line.split()
    ff.append((C.MP(pstr), int(estr)))
  rc = proc.wait()
  if rc: raise ExpectedError, 'factor failed: rc = %d' % rc
  return ff

CMDMAP = {}

def defcommand(f, name = None):
  if isinstance(f, basestring):
    return lambda g: defcommand(g, f)
  else:
    if name is None: name = f.__name__
    CMDMAP[name] = f
    return f

def connect_db(dir):
  db = SQL.connect(OS.path.join(dir, 'db'))
  db.text_factory = str
  c = db.cursor()
  c.executescript(CONNINIT_SQL)
  return db

@defcommand
def setup(dir, kind, groupdesc, gstr, xstr):

  ## Get the group.  This will also figure out the group order.
  G = getgroup(kind, groupdesc)

  ## Figure out the generator order.
  g = G.elt(gstr)
  x = G.elt(xstr)
  ff = []
  m = G.order
  for p, e in factor(m):
    ee = 0
    for i in xrange(e):
      mm = m/p
      t = G.pow(g, mm)
      if not G.idp(t): break
      ee += 1; m = mm
    if ee < e: ff.append((p, e - ee))

  ## Check that x at least has the right order.  This check is imperfect.
  if not G.idp(G.pow(x, m)): raise ValueError, 'x not in <g>'

  ## Prepare the directory.
  try: OS.mkdir(dir)
  except OSError, err: raise ExpectedError, 'mkdir: %s' % err

  ## Prepare the database.
  db = connect_db(dir)
  c = db.cursor()
  c.executescript(SETUP_SQL)

  ## Populate the general information.
  with db:
    c.execute("""INSERT INTO top (kind, groupdesc, g, x, m)
                 VALUES (?, ?, ?, ?, ?)""",
              (kind, groupdesc, G.str(g), G.str(x), str(m)))
    for p, e in ff:
      if p.nbits <= 48: dpbits = 0
      else: dpbits = p.nbits*2/5
      c.execute("""INSERT INTO progress (p, e, dpbits) VALUES (?, ?, ?)""",
                (str(p), e, dpbits))

def get_top(db):
  c = db.cursor()
  c.execute("""SELECT kind, groupdesc, g, x, m, n FROM top""")
  kind, groupdesc, gstr, xstr, mstr, nstr = c.fetchone()
  G = getgroup(kind, groupdesc)
  g, x, m = G.elt(gstr), G.elt(xstr), C.MP(mstr)
  n = nstr is not None and C.MP(nstr) or None
  return G, g, x, m, n

@defcommand
def check(dir):
  rc = [0]
  def bad(msg):
    print >>stderr, '%s: %s' % (PROG, msg)
    rc[0] = 3
  db = connect_db(dir)
  c = db.cursor()
  G, g, x, m, n = get_top(db)
  print '## group: %s %s' % (G.NAME, G.desc)
  print '## g = %s' % G.str(g)
  print '## x = %s' % G.str(x)

  if not G.idp(G.pow(g, m)):
    bad('bad generator/order: %s^%d /= 1' % (G.str(g), m))
  if not G.idp(G.pow(x, m)):
    bad('x not in group: %s^%d /= 1' % (G.str(x), m))

  c.execute("""SELECT p.p, p.e, p.k, p.n, p.dpbits, COUNT(d.z)
               FROM progress AS p LEFT OUTER JOIN points AS d
               ON p.p = d.p AND p.k = d.k
               GROUP BY p.p, p.k
               ORDER BY LENGTH(p.p), p.p""")
  mm = 1
  for pstr, e, k, nnstr, dpbits, ndp in c:
    p, nn = C.MP(pstr), C.MP(nnstr)
    q = p**e
    if m%q:
      bad('incorrect order factorization: %d^%d /| %d' % (p, e, m))
    mm *= q
    if G.idp(G.pow(g, m/p)):
      bad('bad generator/order: %s^{%d/%d} = 1' ^ (G.str(g), m, p))
    r = m/p**k
    h = G.pow(g, r*nn)
    y = G.pow(x, r)
    if not G.eq(h, y):
      bad('bad partial log: (%s^{%d/%d^%d})^%d = %s /= %s = %s^{%d/%d^%d}' %
          (G.str(g), m, p, k, nn, G.str(h), G.str(y), G.str(x), m, p, k))
    if not dpbits or k == e: dpinfo = ''
    else: dpinfo = ' [%d: %d]' % (dpbits, ndp)
    print '## %d: %d/%d%s' % (p, k, e, dpinfo)
  if mm != m:
    bad('incomplete factorization: %d /= %d' % (mm, m))

  if n is not None:
    xx = G.pow(g, n)
    if not G.eq(xx, x):
      bad('incorrect log: %s^%d = %s /= %s' %
          (G.str(g), n, G.str(xx), G.str(x)))
    print '## DONE: %d' % n

  exit(rc[0])

def get_job(db):
  c = db.cursor()
  c.execute("""SELECT p.p, p.e, p.k, p.n, p.dpbits
               FROM progress AS p LEFT OUTER JOIN workers AS w
                       ON p.p = w.p and p.k = w.k
               WHERE p.k < p.e AND (p.dpbits > 0 OR w.pid IS NULL)
               LIMIT 1""")
  row = c.fetchone()
  if row is None: return None, None, None, None, None
  else:
    pstr, e, k, nstr, dpbits = row
    p, n = C.MP(pstr), C.MP(nstr)
    return p, e, k, n, dpbits

@defcommand
def done(dir):
  db = connect_db(dir)
  c = db.cursor()
  G, g, x, m, n = get_top(db)
  if n is not None:
    print '## DONE: %d' % n
    exit(0)
  p, e, k, n, dpbits = get_job(db)
  if p is None: exit(2)
  else: exit(1)

def maybe_cleanup_worker(dir, db, pid):
  c = db.cursor()
  f = OS.path.join(dir, 'lk.%d' % pid)
  state = 'LIVE'
  try: fd = OS.open(f, OS.O_WRONLY)
  except OSError, err:
    if err.errno != E.ENOENT: raise ExpectedError, 'open lockfile: %s' % err
    state = 'STALE'
  else:
    try: F.lockf(fd, F.LOCK_EX | F.LOCK_NB)
    except IOError, err:
      if err.errno != E.EAGAIN: raise ExpectedError, 'check lock: %s' % err
    else:
      state = 'STALE'
  if state == 'STALE':
    try: OS.unlink(f)
    except OSError: pass
    c.execute("""DELETE FROM workers WHERE pid = ?""", (pid,))

def maybe_kill_worker(dir, pid):
  f = OS.path.join(dir, 'lk.%d' % pid)
  try: fd = OS.open(f, OS.O_RDONLY)
  except OSError, err:
    if err.errno != E.ENOENT: raise ExpectedError, 'open lockfile: %s' % err
    return
  try: F.lockf(fd, F.LOCK_EX | F.LOCK_NB)
  except IOError, err:
    if err.errno != E.EAGAIN: raise ExpectedError, 'check lock: %s' % err
  else: return
  OS.kill(pid, SIG.SIGTERM)
  try: OS.unlink(f)
  except OSError: pass

@defcommand
def step(dir, cmd, *args):

  ## Open the database.
  db = connect_db(dir)
  c = db.cursor()
  ##db.isolation_level = 'EXCLUSIVE'

  ## Prepare our lockfile names.
  mypid = OS.getpid()
  nlk = OS.path.join(dir, 'nlk.%d' % mypid)
  lk = OS.path.join(dir, 'lk.%d' % mypid)

  ## Overall exception handling...
  try:

    ## Find out what needs doing and start doing it.  For this, we open a
    ## transaction.
    with db:
      G, g, x, m, n = get_top(db)
      if n is not None: raise ExpectedError, 'job done'

      ## Clear away old workers that aren't doing anything useful any more.
      ## For each worker pid, check that its lockfile is still locked; if
      ## not, it's finished and can be disposed of.
      c.execute("""SELECT pid FROM workers""")
      for pid, in c:
        maybe_cleanup_worker(dir, db, pid)
      for f in OS.listdir(dir):
        if f.startswith('lk.'):
          pid = int(f[3:])
          maybe_cleanup_worker(dir, db, pid)

      ## Find something to do.  Either a job that's small enough for us to
      ## take on alone, and that nobody else has picked up yet, or one that
      ## everyone's pitching in on.
      p, e, k, n, dpbits = get_job(db)
      if p is None: raise ExpectedError, 'no work to do'

      ## Figure out what needs doing.  Let q = p^e, h = g^{m/q}, y = x^{m/q}.
      ## Currently we have n_0 where
      ##
      ##    h^{p^{e-k} n_0} = y^{p^{e-k}}
      ##
      ## Suppose n == n_0 + p^k n' (mod p^{k+1}).  Then p^k n' == n - n_0
      ## (mod p^{k+1}).
      ##
      ##    (h^{p^{e-1}})^{n'} = (g^{m/p})^{n'}
      ##                       = (y/h^{n_0})^{p^{e-k-1}}
      ##
      ## so this is the next discrete log to solve.
      q = p**e
      o = m/q
      h, y = G.pow(g, o), G.pow(x, o)
      hh = G.pow(h, p**(e-1))
      yy = G.pow(G.div(y, G.pow(h, n)), p**(e-k-1))

      ## Take out a lockfile.
      fd = OS.open(nlk, OS.O_WRONLY | OS.O_CREAT, 0700)
      F.lockf(fd, F.LOCK_EX | F.LOCK_NB)
      OS.rename(nlk, lk)

      ## Record that we're working in the database.  This completes our
      ## initial transaction.
      c.execute("""INSERT INTO workers (pid, p, k) VALUES (?, ?, ?)""",
                (mypid, str(p), k))

    ## Before we get too stuck in, check for an easy case.
    if G.idp(yy):
      dpbits = 0 # no need for distinguished points
      nn = 0; ni = 0
    else:

      ## There's nothing else for it.  Start the job up.
      proc = S.Popen([cmd] + list(args) +
                     [str(dpbits), G.NAME, G.desc,
                      G.str(hh), G.str(yy), str(p)],
                     stdin = S.PIPE, stdout = S.PIPE)
      f_in, f_out = proc.stdin.fileno(), proc.stdout.fileno()

      ## Now we must look after it until it starts up.  Feed it stuff on stdin
      ## periodically, so that we notice if our network connectivity is lost.
      ## Collect its stdout.
      for fd in [f_in, f_out]:
        fl = F.fcntl(fd, F.F_GETFL)
        F.fcntl(fd, F.F_SETFL, fl | OS.O_NONBLOCK)
      done = False
      out = ''
      while not done:
        rdy, wry, exy = SEL.select([f_out], [], [], 30.0)
        if rdy:
          while True:
            try: b = OS.read(f_out, 4096)
            except OSError, err:
              if err.errno == E.EAGAIN: break
              else: raise ExpectedError, 'read job: %s' % err
            else:
              if not len(b): done = True; break
              else: out += b
        if not done:
          try: OS.write(f_in, '.')
          except OSError, err: raise ExpectedError, 'write job: %s' % err
      rc = proc.wait()
      if rc: raise ExpectedError, 'job failed: rc = %d' % rc

      ## Parse the answer.  There are two cases.
      if not dpbits:
        nnstr, nistr = out.split()
        nn, ni = C.MP(nnstr), int(nistr)
      else:
        astr, bstr, zstr, nistr = out.split()
        a, b, z, ni = C.MP(astr), C.MP(bstr), G.elt(zstr), int(nistr)

    ## We have an answer.  Start a new transaction while we think about what
    ## this means.
    with db:

      if dpbits:

        ## Check that it's a correct point.
        zz = G.mul(G.pow(hh, a), G.pow(yy, b))
        if not G.eq(zz, z):
          raise ExpectedError, \
              'job incorrect distinguished point: %s^%d %s^%d = %s /= %s' % \
              (hh, a, yy, b, zz, z)

        ## Report this (partial) success.
        print '## [%d, %d/%d: %d]: %d %d -> %s [%d]' % \
            (p, k, e, dpbits, a, b, G.str(z), ni)

        ## If it's already in the database then we have an answer to the
        ## problem.
        c.execute("""SELECT a, b FROM points
                     WHERE p = ? AND k = ? AND z = ?""",
                  (str(p), k, str(z)))
        row = c.fetchone()
        if row is None:
          nn = None
          c.execute("""INSERT INTO points (p, k, a, b, z)
                        VALUES (?, ?, ?, ?, ?)""",
                    (str(p), str(k), str(a), str(b), G.str(z)))
        else:
          aastr, bbstr = row
          aa, bb = C.MP(aastr), C.MP(bbstr)
          if not (b - bb)%p:
            raise ExpectedError, 'duplicate point :-('

          ## We win!
          nn = ((a - aa)*p.modinv(bb - b))%p
          c.execute("""SELECT COUNT(z) FROM points WHERE p = ? AND k = ?""",
                    (str(p), k))
          ni, = c.fetchone()
          print '## [%s, %d/%d: %d] collision %d %d -> %s <- %s %s [#%d]' % \
              (p, k, e, dpbits, a, b, G.str(z), aa, bb, ni)

      ## If we don't have a final answer then we're done.
      if nn is None: return

      ## Check that the log we've recovered is correct.
      yyy = G.pow(hh, nn)
      if not G.eq(yyy, yy):
        raise ExpectedError, 'recovered incorrect log: %s^%d = %s /= %s' % \
            (G.str(hh), nn, G.str(yyy), G.str(yy))

      ## Update the log for this prime power.
      n += nn*p**k
      k += 1

      ## Check that this is also correct.
      yyy = G.pow(h, n*p**(e-k))
      yy = G.pow(y, p**(e-k))
      if not G.eq(yyy, yy):
        raise ExpectedError, 'lifted incorrect log: %s^d = %s /= %s' % \
            (G.str(h), n, G.str(yyy), G.str(yy))

      ## Kill off the other jobs working on this component.  If we crash now,
      ## we lose a bunch of work. :-(
      c.execute("""SELECT pid FROM workers WHERE p = ? AND k = ?""",
                (str(p), k))
      for pid, in c: maybe_kill_worker(dir, pid)
      c.execute("""DELETE FROM workers WHERE p = ? AND k = ?""",
                (str(p), k - 1))
      c.execute("""DELETE FROM points WHERE p = ? AND k = ?""",
                (str(p), k - 1))

      ## Looks like we're good: update the progress table.
      c.execute("""UPDATE progress SET k = ?, n = ? WHERE p = ?""",
                (k, str(n), str(p)))
      print '## [%d, %d/%d]: %d [%d]' % (p, k, e, n, ni)

      ## Quick check: are we done now?
      c.execute("""SELECT p FROM progress WHERE k < e
                   LIMIT 1""")
      row = c.fetchone()
      if row is None:

        ## Wow.  Time to stitch everything together.
        c.execute("""SELECT p, e, n FROM progress""")
        qq, nn = [], []
        for pstr, e, nstr in c:
          p, n = C.MP(pstr), C.MP(nstr)
          qq.append(p**e)
          nn.append(n)
        n = C.MPCRT(qq).solve(nn)

        ## One last check that this is the right answer.
        xx = G.pow(g, n)
        if not G.eq(x, xx):
          raise ExpectedError, \
              'calculated incorrect final log: %s^d = %s /= %s' \
              (G.str(g), n, G.str(xx), G.str(x))

        ## We're good.
        c.execute("""UPDATE top SET n = ?""", (str(n),))
        print '## DONE: %d' % n

  finally:

    ## Delete our lockfile.
    for f in [nlk, lk]:
      try: OS.unlink(f)
      except OSError: pass

    ## Unregister from the database.
    with db:
      c.execute("""DELETE FROM workers WHERE pid = ?""", (mypid,))

PROG = argv[0]

try:
  CMDMAP[argv[1]](*argv[2:])
except ExpectedError, err:
  print >>stderr, '%s: %s' % (PROG, err.message)
  exit(3)