#!/usr/bin/env python3
import json
import os
import socket
import struct


def string(value):
    data = value.encode() if isinstance(value, str) else value
    return struct.pack('<H', len(data)) + data


class Client:
    def __init__(self, address, timeout=5, msize=8192):
        tcp = isinstance(address, tuple)
        family = (socket.AF_INET6 if ':' in address[0] else socket.AF_INET) if tcp else socket.AF_UNIX
        self.socket = socket.socket(family)
        self.socket.settimeout(timeout)
        if tcp:
            self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.tag = 0
        self.fid = 1
        self.msize = msize
        try:
            self.socket.connect(address if tcp else os.fspath(address))
            reply = self.rpc(100, struct.pack('<I', self.msize) + string('9P2000'), tag=65535)
            self.msize = min(self.msize, struct.unpack_from('<I', reply)[0])
            if self.msize < 256 or reply[4:] != string('9P2000'):
                raise OSError('unsupported 9P version or message size')
            attached = self.rpc(104, struct.pack('<II', 1, 0xffffffff) + string('pardes') + string(''))
            if len(attached) != 13:
                raise OSError('invalid attach qid')
        except BaseException:
            self.socket.close()
            raise

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.socket.close()

    def receive(self, count):
        data = bytearray()
        while len(data) < count:
            chunk = self.socket.recv(count - len(data))
            if not chunk:
                raise EOFError('9P peer closed the connection')
            data.extend(chunk)
        return bytes(data)

    def rpc(self, kind, payload, tag=None):
        if tag is None:
            self.tag = (self.tag + 1) % 65535
            tag = self.tag
        packet = struct.pack('<IBH', 7 + len(payload), kind, tag) + payload
        if len(packet) > self.msize:
            raise ValueError('request exceeds negotiated message size')
        self.socket.sendall(packet)
        size, answer, received_tag = struct.unpack('<IBH', self.receive(7))
        if not 7 <= size <= self.msize or received_tag != tag:
            raise OSError('invalid 9P reply header')
        reply = self.receive(size - 7)
        if answer == 107:
            length = struct.unpack_from('<H', reply)[0]
            raise OSError(reply[2:2 + length].decode(errors='replace'))
        if answer != kind + 1:
            raise OSError(f'expected reply {kind + 1}, received {answer}')
        return reply

    def open(self, path, mode=0):
        names = [part for part in os.fspath(path).split('/') if part]
        self.fid += 1
        fid = self.fid
        oldfid = 1
        try:
            for start in range(0, max(1, len(names)), 16):
                part = names[start:start + 16]
                reply = self.rpc(110, struct.pack('<IIH', oldfid, fid, len(part)) +
                                 b''.join(string(name) for name in part))
                walked = struct.unpack_from('<H', reply)[0]
                if len(reply) != 2 + 13 * walked:
                    raise OSError('invalid walk qids')
                if walked != len(part):
                    raise FileNotFoundError(path)
                oldfid = fid
            if len(self.rpc(112, struct.pack('<IB', fid, mode))) != 17:
                raise OSError('invalid open reply')
            return fid
        except BaseException:
            try:
                self.close(fid)
            except OSError:
                pass
            raise

    def close(self, fid):
        self.rpc(120, struct.pack('<I', fid))

    def read_fid(self, fid, offset=0, count=None):
        count = min(count if count is not None else self.msize - 11, self.msize - 11)
        reply = self.rpc(116, struct.pack('<IQI', fid, offset, count))
        size = struct.unpack_from('<I', reply)[0]
        if size != len(reply) - 4 or size > count:
            raise OSError('invalid 9P read length')
        return reply[4:]

    def read(self, path):
        fid = self.open(path)
        try:
            data = bytearray()
            while chunk := self.read_fid(fid, len(data)):
                data.extend(chunk)
            return bytes(data)
        finally:
            self.close(fid)

    def screen(self):
        return json.loads(self.read('/self/screen'))

    def write(self, path, data, truncate=False):
        fid = self.open(path, 1 | (16 if truncate else 0))
        try:
            offset = 0
            while offset < len(data):
                chunk = data[offset:offset + self.msize - 23]
                reply = self.rpc(118, struct.pack('<IQI', fid, offset, len(chunk)) + chunk)
                count = struct.unpack('<I', reply)[0]
                if not 0 < count <= len(chunk):
                    raise OSError('invalid 9P write length')
                offset += count
        finally:
            self.close(fid)

    def list(self, path='/'):
        data = self.read(path)
        names = []
        offset = 0
        while offset < len(data):
            size = struct.unpack_from('<H', data, offset)[0]
            if size < 47 or offset + size + 2 > len(data):
                raise OSError('invalid 9P directory entry')
            length = struct.unpack_from('<H', data, offset + 41)[0]
            if 43 + length > size + 2:
                raise OSError('invalid directory name length')
            names.append(data[offset + 43:offset + 43 + length].decode())
            offset += size + 2
        return names