#!/usr/bin/env python

# System modules
import time
import struct
import serial
from astropy.io import fits
from astropy.logger import log
import numpy as np

KONTINUE = b'K'
RESEND = b'R'
STOP = b'S'
BAUDRATE = 115200

class AllSky340:
    def __init__(self, device, timeout=0.160):
        self.ser = serial.serial_for_url(device)
        self.ser.baudrate = BAUDRATE
        self.ser.bytesize = serial.EIGHTBITS
        self.ser.timeout = timeout
        for i in range(3):
            self.ping()

    def wait(self, nbytes):
        '''
        Calculate the required timeout to transmit a certain number of bytes
        based on the current baud rate of the serial port

        ser -- the serial.Serial() to receive from
        nbytes -- the number of bytes that will be transmitted

        return -- the time required to transmit in seconds
        '''
        return (BAUDRATE / (nbytes * 8)) * 1.5


    def checksum(self, c):
        inv = ~c[0] & 0xFF
        mask = ~(1 << 7)
        checksum = inv & mask
        if len(c) > 1:
            for i in range(1, len(c)):
                inv = ~c[i] & 0xFF
                xor = checksum ^ inv
                checksum = xor & mask
        log.info(f"Checksum: {chr(checksum).encode()}")
        return chr(checksum).encode()


    def command(self, cmd, nbytes=None):

        check = self.checksum(cmd)
        to_send = cmd + check
        log.info(f"Sending:  {to_send}")
        self.ser.write(to_send)
        time.sleep(0.1)
        resp = self.ser.read(1)
        if resp==check:
            log.info(f"Response ok: '{resp}' == '{check}'")
        else:
            log.warning(f"Wired response: '{resp}' == '{check}'")

        if nbytes:
            resp = self.ser.read(nbytes)
            log.info(f"Received: {resp}")
            return resp
        else:
            return True


    def ping(self):
        log.info("Pinging Camera")
        resp = self.command(b"E", 1)
        if resp == b"O" :
            log.info(f"Response ok: {resp}")
            return True
        else:
            log.warning(f"Response not ok: {resp}")
            return False


    def open_shutter(self):
        log.info("Opening camera shutter")
        return self.command(b"O")


    def close_shutter(self):
        log.info("Closing camera shutter")
        return self.command(b"C")


    def shutter_off(self):
        log.info("De-energizing camera shutter")
        self.command(b"K")
        return self.command(b"K")


    def abort(self):
        log.info("Aborting current exposure")
        return self.command(b"A")


    def firmware(self):
        resp = self.command(b"V", 2)
        minor = resp[1]
        major = resp[0]
        if major >= 8:
            ver = "T"
            major -= 8
        else:
            ver = "V"
        version = f"{ver}{major}{minor}"
        log.info(f"Camera firmware: {version}")
        return version


    def serial(self):
        resp = self.command(b"r", 9)
        version = resp.decode("utf-8").strip("\r")
        log.info(f"Camera serial: {version}")
        return version


    def get_min_move(self):
        resp = self.command(b"n", 2)
        log.info(f"Camera min_move: {resp[0]} {resp[1]}")
        return resp


    def set_min_move(self):
        data = b"N"  + struct.pack("<2B", *[17,0])
        resp = self.command(data, 3)
        log.info(f"Camera min_move: {resp}")
        return resp


    def heater_on(self):
        cmd = b'g' + struct.pack("B", 1)
        log.info("Turning camera heater on")
        self.command(cmd)
        return True


    def heater_off(self):
        cmd = b'g' + struct.pack("B", 0)
        log.info("Turning camera heater off")
        self.command(cmd)
        return True


    def shutter_chop_on(self):
        cmd = b'U' + struct.pack("B", 1)
        log.info("Enabling chopping shutter mode.")
        self.command(cmd)
        return True


    def shutter_chop_off(self):
        cmd = b'U' + struct.pack("B", 0)
        log.info("Disabling chopping shutter mode.")
        self.command(cmd)
        return True


    def take_image(self, exposure, light=True):

        exptime = exposure * 10000.0
        if exptime > 0x63FFFF:
            exptime = 0x63FFFF;

        #exp = int(exptime).to_bytes(3, byteorder='little')
        exp = struct.pack("<I", int(exptime))[:3]

        combuf = bytearray(6)
        combuf[0] = ord("T")
        combuf[1] = exp[2]
        combuf[2] = exp[1]
        combuf[3] = exp[0]
        combuf[4] = 0x00
        combuf[5] = 0x01 if light else 0x00

        log.info("Starting exposure")
        self.command(combuf)

        state = {
            b"E" : "Expose",
            b"R" : "Readout",
            b""  : "Reading",
            b"D" : "Done",
        }

        out = b"E"
        while out != b"D":
            time.sleep(0.1)
            out = self.ser.read(1)
            try:
                log.info(state[out])
            except:
                log.warning(f"Wired state not in dict: {out}")



    def transfer_block(self, chunk_size=4096, ignore_cs=False, tries=3):

        data = b""
        for i in range(tries):
            log.info(f"Get Image Block: try {i}")

            # not the first try, transmit checksum error so the camera will try again
            if i > 0:
                log.warning("Resending")
                self.ser.write(RESEND)

            # data = b''
            # while len(data) < chunk_size*2  :
            #     #time.sleep(0.05)
            #     res = self.ser.read(8)
            #     print(f"reading....{len(data)} : {res}")
            #     data += res

            data = self.ser.read(chunk_size*2)
            self.wait(chunk_size)

            csum_byte = self.ser.read(1)
            try:
                csum_byte = ord(csum_byte)
            except TypeError as e:
                log.error(f"Empty csum_byte: {csum_byte}")
                return data

            # not enough bytes, therefore transfer failed
            if len(data) != chunk_size*2 :
                log.warning('Not enough data returned before timeout')
            else:
                log.info('Data lenght ok')

            csum = 0
            for byte in data:
                csum ^= (byte & 0xff)
                csum &= 0xff
                # if csum == csum_byte:
                #     print(f"{byte} {csum :>4} {csum_byte}")

            if csum == csum_byte or ignore_cs:
                log.info('Kontinue')
                self.ser.write(KONTINUE)
                return data

            log.warning("Checksum Error")
            #valid = True
            #else:
            #    cs_failed += 1
            #    log.warning('Resend')
            #    self.command(RESEND)

            #if cs_failed > 0:
            #    log.error('Checksum failed')
            #else:
            #    log.info('Checksum ok')

        # too many retries passed, abort
        log.error('Get Image Block: retries exhausted, abort transfer')
        self.command(STOP)
        return data

    def create_fits(self):

        # Download Image
        blocks_expected = int((640 * 480) / 4096)

        # Download Image

        data = b""
        blocks_complete = 0

        for _ in range(blocks_expected):
            blocks_complete += 1

            self.command(b"X")
            time.sleep(1)
            
            data += self.transfer_block()
            time.sleep(1)
            
            print(f"{blocks_complete} or {_} of {blocks_expected}")
            
        # head = fits.Header()

        # # Now make into a fits image
        # #data = np.fromstring(data, dtype=np.uint16)
        # data = np.frombuffer(data, dtype=np.uint16)
        # data = data.reshape((480, 640))

        return data


if __name__ == '__main__':
#    c = AllSky340(device="socket://151.3.45.60:8899")
    c = AllSky340(device="/dev/ttyr00")
    c.abort()
    c.open_shutter()
    c.take_image(1.2345)
    c.close_shutter()
    c.shutter_off()

    time.sleep(1)

    # Download Image
    blocks_expected = int((640 * 480) / 4096)

    data = b''
    blocks_complete = 0

    for _ in range(blocks_expected):
        data += c.transfer_image()
        blocks_complete += 1
        print(f"{blocks_complete} or {_} of {blocks_expected}")

    head = fits.Header()

    # Now make into a fits image
    #data = np.fromstring(data, dtype=np.uint16)
    data = np.frombuffer(data, dtype=np.uint16)
    data = data.reshape((480, 640))

    hdu = fits.PrimaryHDU(data, header=head)
    hdu.write("ciccio.fits")
