import sys
import os
import re
import time
import struct
import logging
import zlib
import bz2

from libcan.filter import COMPRESS_GZIP, COMPRESS_GZIP_ALL, COMPRESS_BZIP2, COMPRESS_BZIP2_ALL, COMPRESS_NONE, COMPRESS_CONST, ENCRYPT_CONST, TempFile, ChecksumWriter, BlowfishWriter, GzipWriter, BZ2Writer, GzipReader, BZ2Reader, BlowfishReader, PlainReader
from libcan.file import *
from libcan import __version__, BUFSIZE


writer_tmpl = "%Y%m%d-%H%M"
reader_tmpl = re.sub(r"%[Y]", r"\d{4}", writer_tmpl)
reader_tmpl = re.sub(r"%[mdHMS]", r"\d{2}", reader_tmpl)


class SnapshotError(Exception):
    pass


class SnapshotCollection(object):
    # FIXME limit open files to a certain number.

    def __init__(self, source, keyfile):
        self.source = source
        self.keyfile = keyfile

    @classmethod
    def from_name(cls, source, keyfile=None, ignore_corrupt=False):
        obj = cls(source, keyfile)

        directory, name = os.path.split(source)
        directory = directory or "."

        obj.snapshots = {}

        # If there are backup snapshots already, examine them and
        # use their data.
        if os.path.exists(directory):
            # Collect snapshots.
            bad_snapshots = [os.path.join(directory, n) for n in os.listdir(directory)
                             if re.match("^" + name + "-" + reader_tmpl + r"\.tmp\.\d+$", n)]
            bad_snapshots.sort()

            # Remove all bad/incomplete snapshots.
            while bad_snapshots:
                snapshot = bad_snapshots.pop()

                # Test if another can process is still running.
                pid = snapshot.rsplit(".", 1)[-1]
                pid = int(pid)
                try:
                    os.kill(pid, 0)
                except OSError:
                    if ignore_corrupt:
                        logging.warn("ignore corrupt snapshot %s", snapshot)
                    else:
                        logging.warn("remove corrupt snapshot %s", snapshot)
                        os.remove(snapshot)
                else:
                    logging.error("snapshot %s in progress", snapshot)
                    sys.exit(1)

            snapshots = [os.path.join(directory, n) for n in os.listdir(directory)
                         if re.match("^" + name + "-" + reader_tmpl + "$", n)]
            snapshots.sort()

            for filename in snapshots:
                fileobj = open(filename)
                serial = SnapshotReader.read_header(fileobj)[3]
                fileobj.close()
                obj.snapshots[serial] = filename
        return obj

    @classmethod
    def from_snapshot(cls, filename, keyfile=None, ignore_corrupt=False):
        fileobj = open(filename)
        source = SnapshotReader.read_header(fileobj)[0]
        fileobj.close()
        return cls.from_name(source, keyfile, ignore_corrupt)

    def close(self):
        for fileobj in self.snapshots.itervalues():
            if hasattr(fileobj, "close"):
                fileobj.close()

    def get_snapshot_by_serial(self, serial):
        snapshot = self.snapshots[serial]
        if not isinstance(snapshot, SnapshotReader):
            snapshot = SnapshotReader(snapshot, self, self.keyfile)
            self.snapshots[snapshot.serial] = snapshot
        return snapshot

    def get_snapshot_by_name(self, name):
        for serial, snapshot in self.snapshots.iteritems():
            if isinstance(snapshot, SnapshotReader):
                snapshot = snapshot.filename
            if os.path.basename(name) == os.path.basename(snapshot):
                return self.get_snapshot_by_serial(serial)
        else:
            raise KeyError

    def get_last_snapshot(self):
        if self.snapshots:
            serials = self.snapshots.keys()
            serials.sort()
            return self.get_snapshot_by_serial(serials[-1])
        return None


class SnapshotReader(object):

    @staticmethod
    def read_header(fileobj):
        if fileobj.read(4) != "CAN4":
            raise SnapshotError("no CAN4 archive")
        try:
            source = unpack_string(fileobj)
            if not source:
                raise SnapshotError("empty source field")
            created = unpack_string(fileobj)
            timestamp = unpack_large(fileobj)
            serial = unpack_large(fileobj)
            compression = unpack_small(fileobj)
            if compression >= COMPRESS_CONST:
                raise SnapshotError("invalid compression method %r" % compression)
            encryption = unpack_small(fileobj)
            if encryption >= ENCRYPT_CONST:
                raise SnapshotError("invalid encryption method %r" % encryption)
        except TypeError:
            # Raised if ord() gets a zero-length string.
            raise SnapshotError("damaged CAN4 archive")
        return source, created, timestamp, serial, compression, encryption

    def __init__(self, filename, collection, keyfile=None):
        self.filename = filename
        self.collection = collection
        self.keyfile = keyfile

        self.fileobj = open(filename)
        self.source, self.created, self.timestamp, self.serial, self.compression, self.encryption = self.read_header(self.fileobj)
        logging.noise("CAN4 archive: %s, serial no.: %d", os.path.basename(self.filename), self.serial)

        if self.encryption:
            if self.keyfile is None:
                raise SystemExit("snapshot is encrypted and no keyfile was specified!")
            password = open(self.keyfile).read().strip()
            self.fileobj = BlowfishReader(self.fileobj, password)

        if self.compression in (COMPRESS_NONE, COMPRESS_GZIP, COMPRESS_BZIP2):
            self.fileobj = PlainReader(self.fileobj, None)
        elif self.compression == COMPRESS_GZIP_ALL:
            self.fileobj = GzipReader(self.fileobj, None)
        elif self.compression == COMPRESS_BZIP2_ALL:
            self.fileobj = BZ2Reader(self.fileobj, None)
        else:
            raise SnapshotError("invalid compression method %r" % self.compression)

        self.next_offset = self.fileobj.tell()
        assert self.next_offset == 0

    def __iter__(self):
        while True:
            try:
                file = self.next()
            except EOFError:
                break

            if file.isref():
                file = self.dereference(file)
            yield file

    def next(self):
        self.fileobj.seek(self.next_offset)

        offset = self.fileobj.tell()
        type = unpack_small(self.fileobj)
        #logging.noise("file type = %s", {0: "file", 1: "ref", 255: "eos"}[type])

        if type == TYPE_EOF:
            # End of file marker.
            raise EOFError

        elif type == TYPE_FILE:
            # File marker. Full metadata and data.
            file = File.from_snapshot(self.fileobj)
            file.snapshot = self.serial
            file.offset = offset
            if file.isreg():
                # Jump data part.
                file.compression = unpack_small(self.fileobj)
                file.compressed_size = unpack_large(self.fileobj)
                file.checksum = unpack_large(self.fileobj)
                file.data_offset = self.fileobj.tell()
                self.next_offset = self.fileobj.tell() + file.compressed_size
            else:
                self.next_offset = self.fileobj.tell()

        elif type == TYPE_REFERENCE:
            # Full reference to a file.
            file = Reference.from_snapshot(self.fileobj)
            self.next_offset = self.fileobj.tell()

        else:
            if self.keyfile:
                raise SnapshotError("invalid data, wrong keyfile?")
            else:
                raise SnapshotError("invalid file type value %r at 0x%06x" % (type, offset))

        return file

    def get_file(self, offset, serial):
        self.next_offset = offset

        try:
            file = self.next()
        except EOFError:
            raise SnapshotError("file 0x%06x-%d not found" % (offset, serial))

        assert file.serial == serial
        return file

    def close(self):
        self.fileobj.close()

    def dereference(self, ref):
        assert isinstance(ref, Reference)
        snapshot = self.collection.get_snapshot_by_serial(ref.snapshot)
        file = snapshot.get_file(ref.offset, ref.target)
        assert file.offset == ref.offset
        assert file.serial == ref.target
        assert file.snapshot != self.serial
        assert not isinstance(file, Reference)
        return file


class SnapshotWriter(object):

    def __init__(self, source, root, serial, compression, keyfile):
        self.source = source
        self.directory, self.name = os.path.split(source or "snapshot")
        self.directory = self.directory or "."

        self.root = root
        self.serial = serial
        self.compression = compression
        self.keyfile = keyfile

        self.filename = os.path.join(self.directory, time.strftime(self.name + "-" + writer_tmpl))

        # If the database file itself exists, we abort.
        if os.path.exists(self.filename):
            raise SnapshotError("snapshot %s exists!" % self.filename)

        self.file_serial = 0
        self.level = 9

        logging.info("create snapshot: %s", self.filename)

        self.fileobj = open(self.filename + ".tmp.%d" % os.getpid(), "w")
        self.write_header()

        if keyfile is not None:
            password = open(keyfile).read().strip()
            self.fileobj = BlowfishWriter(self.fileobj, password)

        if self.compression == COMPRESS_GZIP_ALL:
            self.fileobj = GzipWriter(self.fileobj)
        elif self.compression == COMPRESS_BZIP2_ALL:
            self.fileobj = BZ2Writer(self.fileobj)

    def add(self, file):
        """Write a record to the snapshot file.
        """
        if file.is_deleted():
            return

        elif file.is_new() or file.is_modified():
            # add the file to the snapshot
            if file.isreg():
                fileobj = open(os.path.join(self.root, file.name))
                self.write_file(file, fileobj)
                fileobj.close()
            else:
                self.write_file(file)

        elif file.is_unchanged():
            self.write_reference(file)

        self.file_serial += 1

    def close(self):
        """Close the snapshot file.
        """
        self.fileobj.write(pack_small(TYPE_EOF))
        self.fileobj.close()
        os.rename(self.filename + ".tmp.%d" % os.getpid(), self.filename)

    # --- writer methods
    def write_header(self):
        self.fileobj.write("CAN4")
        self.fileobj.write(pack_string(self.source))
        self.fileobj.write(pack_string(__version__))
        self.fileobj.write(pack_large(int(time.time())))
        self.fileobj.write(pack_large(self.serial))
        self.fileobj.write(pack_small(self.compression))
        self.fileobj.write(pack_small(self.keyfile is not None))

    def write_file(self, file, fileobj=None):
        if file.isreg():
            assert fileobj is not None
        else:
            assert fileobj is None

        self.fileobj.write(pack_small(TYPE_FILE))

        if fileobj is not None:
            compression = self.compression
            if file.is_compressed():
                compression = COMPRESS_NONE
            tempfile, file.size, compressed_size, checksum = self.prepare_data(compression, fileobj)
            self.write_metadata(file)
            self.write_data(tempfile, compression, compressed_size, checksum)
        else:
            self.write_metadata(file)

    def write_metadata(self, file):
        self.fileobj.write(pack_large(self.file_serial))
        file.write_metadata(self.fileobj)

    def prepare_data(self, compression, fileobj):
        """Copy the data from fileobj to a temporary file object
           compressing it on demand and calculating its checksum.
        """
        tempfile = TempFile()
        if compression == COMPRESS_GZIP:
            filter = GzipWriter(tempfile)
        elif compression == COMPRESS_BZIP2:
            filter = BZ2Writer(tempfile)
        else:
            filter = tempfile
        filter = ChecksumWriter(filter)

        real_size = 0
        while True:
            buf = fileobj.read(BUFSIZE)
            if not buf:
                break
            real_size += len(buf)
            filter.write(buf)
        filter.finish()

        compressed_size = tempfile.tell()
        checksum = filter.get_checksum()
        tempfile.rewind()
        return tempfile, real_size, compressed_size, checksum

    def write_data(self, tempfile, compression, compressed_size, checksum):
        self.fileobj.write(pack_small(compression))
        self.fileobj.write(pack_large(compressed_size))
        self.fileobj.write(pack_large(checksum))

        while compressed_size > 0:
            buf = tempfile.read(BUFSIZE)
            if not buf:
                raise ValueError("what's this?")
            compressed_size -= len(buf)
            self.fileobj.write(buf)
        tempfile.close()

    def write_reference(self, file):
        reference = Reference.from_file(file)
        self.fileobj.write(pack_small(TYPE_REFERENCE))
        self.fileobj.write(pack_large(self.file_serial))
        reference.write_reference(self.fileobj)


