# Copyright (C) 2009 Kris Maglione

import operator
import os
import re
import sys
from threading import *
import traceback

import pyxp
from pyxp import fcall, fields
from pyxp.mux import Mux
from pyxp.types import *

if os.environ.get('NAMESPACE', None):
    namespace = os.environ['NAMESPACE']
else:
    try:
        namespace = '/tmp/ns.%s.%s' % (
            os.environ['USER'], 
            re.sub(r'\.0$', '', os.environ['DISPLAY']))
    except Exception:
        pass
NAMESPACE = namespace

OREAD = 0x00
OWRITE = 0x01
ORDWR = 0x02
OEXEC = 0x03
OEXCL = 0x04
OTRUNC = 0x10
OREXEC = 0x20
ORCLOSE = 0x40
OAPPEND = 0x80

ROOT_FID = 0

class ProtocolException(Exception):
    pass
class RPCError(Exception):
    pass

class Client(object):
    ROOT_FID = 0

    @staticmethod
    def respond(callback, data, exc=None, tb=None):
        if hasattr(callback, 'func_code'):
            callback(*(data, exc, tb)[0:callback.func_code.co_argcount])
        elif callable(callback):
            callback(data)

    def __enter__(self):
        return self
    def __exit__(self, *args):
        self._cleanup()

    def __init__(self, conn=None, namespace=None, root=None):
        if not conn and namespace:
            conn = 'unix!%s/%s' % (NAMESPACE, namespace)
        try:
            self.lastfid = ROOT_FID
            self.fids = set()
            self.lock = RLock()

            def process(data):
                return fcall.Fcall.unmarshall(data)[1]
            self.mux = Mux(conn, process, maxtag=256)

            resp = self._dorpc(fcall.Tversion(version=pyxp.VERSION, msize=65535))
            if resp.version != pyxp.VERSION:
                raise ProtocolException, "Can't speak 9P version '%s'" % resp.version
            self.msize = resp.msize

            self._dorpc(fcall.Tattach(fid=ROOT_FID, afid=fcall.NO_FID,
                       uname=os.environ['USER'], aname=''))

            if root:
                path = self._splitpath(root)
                resp = self._dorpc(fcall.Twalk(fid=ROOT_FID,
                                               newfid=ROOT_FID,
                                               wname=path))
        except Exception:
            traceback.print_exc(sys.stdout)
            if getattr(self, 'mux', None):
                self.mux.fd.close()
            raise

    def _cleanup(self):
        try:
            for f in self.files:
                f.close()
        finally:
            self.mux.fd.close()
            self.mux = None

    def _dorpc(self, req, callback=None, error=None):
        def doresp(resp):
            if isinstance(resp, fcall.Rerror):
                raise RPCError, "%s[%d] RPC returned error: %s" % (
                    req.__class__.__name__, resp.tag, resp.ename)
            if req.type != resp.type ^ 1:
                raise ProtocolException, "Missmatched RPC message types: %s => %s" % (
                    req.__class__.__name__, resp.__class__.__name__)
            return resp

        def next(mux, resp):
            try:
                res = doresp(resp)
            except Exception, e:
                self.respond(error or callback, None, e, None)
            else:
                self.respond(callback, res)

        if not callback:
            return doresp(self.mux.rpc(req))
        self.mux.rpc(req, next)

    def _splitpath(self, path):
        if isinstance(path, list):
            return path
        return [v for v in path.split('/') if v != '']

    def _getfid(self):
        with self.lock:
            if self.fids:
                return self.fids.pop()
            self.lastfid += 1
            return self.lastfid
    def _putfid(self, fid):
        with self.lock:
            self.fids.add(fid)

    def _aclunk(self, fid, callback=None):
        def next(resp, exc, tb):
            if resp:
                self._putfid(fid)
            self.respond(callback, resp, exc, tb)
        self._dorpc(fcall.Tclunk(fid=fid), next)

    def _clunk(self, fid):
        try:
            self._dorpc(fcall.Tclunk(fid=fid))
        finally:
            self._putfid(fid)

    def _walk(self, path):
        fid = self._getfid()
        ofid = ROOT_FID
        while True:
            self._dorpc(fcall.Twalk(fid=ofid, newfid=fid,
                                   wname=path[0:fcall.MAX_WELEM]))
            path = path[fcall.MAX_WELEM:]
            ofid = fid
            if len(path) == 0:
                break

        @apply
        class Res:
            def __enter__(res):
                return fid
            def __exit__(res, exc_type, exc_value, traceback):
                if exc_type:
                    self._clunk(fid)
        return Res

    _file = property(lambda self: File)
    def _open(self, path, mode, fcall, origpath=None):
        resp = None

        with self._walk(path) as nfid:
            fid = nfid
            fcall.fid = fid
            resp = self._dorpc(fcall)

        def cleanup():
            self._aclunk(fid)
        file = self._file(self, origpath or '/'.join(path), resp, fid, mode, cleanup)
        return file

    def open(self, path, mode=OREAD):
        path = self._splitpath(path)

        return self._open(path, mode, fcall.Topen(mode=mode))

    def create(self, path, mode=OREAD, perm=0):
        path = self._splitpath(path)
        name = path.pop()

        return self._open(path, mode, fcall.Tcreate(mode=mode, name=name, perm=perm),
                          origpath='/'.join(path + [name]))

    def remove(self, path):
        path = self._splitpath(path)

        with self._walk(path) as fid:
            self._dorpc(fcall.Tremove(fid=fid))

    def stat(self, path):
        path = self._splitpath(path)

        try:
            with self._walk(path) as fid:
                resp = self._dorpc(fcall.Tstat(fid= fid))
                st = resp.stat
                self._clunk(fid)
            return st
        except RPCError:
            return None

    def read(self, path, *args, **kwargs):
        with self.open(path) as f:
            return f.read(*args, **kwargs)
    def readlines(self, path, *args, **kwargs):
        with self.open(path) as f:
            for l in f.readlines(*args, **kwargs):
                yield l
    def readdir(self, path, *args, **kwargs):
        with self.open(path) as f:
            for s in f.readdir(*args, **kwargs):
                yield s
    def write(self, path, *args, **kwargs):
        with self.open(path, OWRITE) as f:
            return f.write(*args, **kwargs)

class File(object):

    def __enter__(self):
        return self
    def __exit__(self, *args):
        self.close()

    def __init__(self, client, path, fcall, fid, mode, cleanup):
        self.lock = RLock()
        self.client = client
        self.path = path
        self.fid = fid
        self._cleanup = cleanup
        self.mode = mode
        self.iounit = fcall.iounit
        self.qid = fcall.qid
        self.closed = False

        self.offset = 0
    def __del__(self):
        if not self.closed:
            self._cleanup()

    def _dorpc(self, fcall, async=None, error=None):
        if hasattr(fcall, 'fid'):
            fcall.fid = self.fid
        return self.client._dorpc(fcall, async, error)

    def stat(self):
        resp = self._dorpc(fcall.Tstat())
        return resp.stat

    def read(self, count=None, offset=None, buf=''):
        if count is None:
            count = self.iounit
        res = []
        with self.lock:
            offs = self.offset
            if offset is not None:
                offs = offset
            while count > 0:
                n = min(count, self.iounit)
                count -= n

                resp = self._dorpc(fcall.Tread(offset=offs, count=n))
                data = resp.data

                offs += len(data)
                res.append(data)

                if len(data) < n:
                    break
            if offset is None:
                self.offset = offs
        return ''.join(res)

    def readlines(self):
        last = None
        while True:
            data = self.read()
            if not data:
                break
            lines = data.split('\n')
            if last:
                lines[0] = last + lines[0]
                last = None
            for i in range(0, len(lines) - 1):
                yield lines[i]
            last = lines[-1]
        if last:
            yield last

    def write(self, data, offset=None):
        if offset is None:
            offset = self.offset
        off = 0
        with self.lock:
            offs = self.offset
            if offset is not None:
                offs = offset
            while off < len(data):
                n = min(len(data), self.iounit)

                resp = self._dorpc(fcall.Twrite(offset=offs,
                                               data=data[off:off+n]))
                off += resp.count
                offs += resp.count
                if resp.count < n:
                    break
            if offset is None:
                self.offset = offs
        return off
    def readdir(self):
        if not self.qid.type & Qid.QTDIR:
            raise Exception, "Can only call readdir on a directory"
        off = 0
        while True:
            data = self.read(self.iounit, off)
            if not data:
                break
            off += len(data)
            for s in Stat.unmarshall_list(data):
                yield s

    def close(self):
        assert not self.closed
        self.closed = True
        try:
            self._cleanup()
        except:
            pass
        self.tg = None
        self.fid = None
        self.client = None
        self.qid = None

    def remove(self):
        try:
            self._dorpc(fcall.Tremove())
        finally:
            self.close()

# vim:se sts=4 sw=4 et: