#!/usr/bin/env python
#
# Copyright 2017 Ettus Research (National Instruments)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
"""
N310 Built-In Self Test (BIST)
"""

from __future__ import print_function
import sys
import socket
import json
from datetime import datetime
import argparse
import pyudev

def post_results(results):
    """
    Given a dictionary, post the results.

    This will print the results as JSON to stdout.
    """
    print(json.dumps(
        results,
        sort_keys=True,
        indent=4,
        separators=(',', ': ')
    ))


class N310BIST(object):
    """
    BIST Tool for the USRP N3xx series
    """
    # This defines special tests that are really collections of other tests.
    collections = {
        'standard': ["ddr3", "gpsdo", "rtc", "temp", "clock_int", "tpm"],
        'extended': "*",
    }

    @staticmethod
    def make_arg_parser():
        """
        Return arg parser
        """
        parser = argparse.ArgumentParser(
            description="N3xx BIST Tool",
        )
        parser.add_argument(
            '-n', '--dry-run', action='store_true',
            help="Fake out the tests. All tests will return a valid response," \
                 " but will not actually interact with hardware.",
        )
        parser.add_argument(
            'tests',
            help="List the tests that should be run",
            nargs='+', # There has to be at least one
        )
        return parser

    def __init__(self):
        self.args = N310BIST.make_arg_parser().parse_args()
        self.tests_to_run = set()
        for test in self.args.tests:
            if test in self.collections:
                for test in self.expand_collection(test):
                    self.tests_to_run.add(test)
            else:
                self.tests_to_run.add(test)

    def expand_collection(self, coll):
        """
        Return names of tests in a collection
        """
        tests = self.collections[coll]
        if tests == "*":
            tests = {x.replace('bist_', '')
                     for x in dir(self)
                     if x.find('bist_') == 0
                    }
        else:
            tests = set(tests)
        return tests

    def run(self):
        """
        Execute tests.

        Returns True on Success.
        """
        def execute_test(testname):
            """
            Actually run a test.
            """
            testmethod_name = "bist_{0}".format(testname)
            sys.stderr.write(
                "Executing test method: {0}\n\n".format(testmethod_name)
            )
            try:
                return getattr(self, testmethod_name)()
            except AttributeError:
                sys.stderr.write("Test not defined: {}\n".format(testname))
                return False, {}
        tests_successful = True
        result = {}
        for test in self.tests_to_run:
            status, result_data = execute_test(test)
            tests_successful = tests_successful and status
            if tests_successful:
                result[test] = result_data
        post_results(result)
        return tests_successful

#############################################################################
# BISTS
# All bist_* methods must return True/False success values!
#############################################################################
    def bist_rtc(self):
        """
        BIST for RTC (real time clock)

        Return dictionary:
        - time: Returns the current UTC time, with seconds-accuracy, in ISO 8601
                format, as a string.

        Return status:
        Unless the 'date' command fails (which is used under the hood), will
        always return True.
        """
        assert 'rtc' in self.tests_to_run
        result = {}
        try:
            utc_time = datetime.utcnow().replace(microseconds=0).isoformat()
            result['time'] = utc_time + "+00:00"
        except:
            if self.args.dry_run:
                result['time'] = '2017-07-11T17:31:03+00:00'
            else:
                return False, {}
        return True, result

    def bist_ddr3(self):
        """
        BIST for PL DDR3 DRAM
        Description: Calls a test to examine the speed of the DDR3
        External Equipment: None
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'speed': Return numeric value of speed in MB/s
            }
        """
        assert 'ddr3' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_gpsdo(self):
        """
        BIST for GPSDO
        Description: Returns GPS information
        External Equipment: None; Recommend attaching an antenna or providing
                           fake GPS information

        Return dictionary: A TPV dictionary as returned by gpsd.
        See also: http://www.catb.org/gpsd/gpsd_json.html

        Check for mode 2 or 3 to see if it's locked.
        """
        assert 'gpsdo' in self.tests_to_run
        if self.args.dry_run:
            return True, {
                "class": "TPV",
                "time": "2017-04-30T11:48:20.10Z",
                "ept": 0.005,
                "lat": 30.407899,
                "lon": -97.726634,
                "alt": 1327.689,
                "epx": 15.319,
                "epy": 17.054,
                "epv": 124.484,
                "track": 10.3797,
                "speed": 0.091,
                "climb": -0.085,
                "eps": 34.11,
                "mode": 3
            }
        my_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        my_sock.connect(('localhost', 2947))
        sys.stderr.write("Connected to GPSDO socket.\n")
        query_cmd = b'?WATCH={"enable":true,"json":true}'
        my_sock.sendall(query_cmd)
        sys.stderr.write("Sent query: {}\n".format(query_cmd))
        while my_sock.recv(1) != b'\n':
            pass
        sys.stderr.write("Received initial newline.\n")
        result = {}
        while result.get('class', None) != 'TPV':
            json_result = b''
            next_char = b''
            while next_char != b'\n':
                json_result += next_char
                next_char = my_sock.recv(1)
            sys.stderr.write(
                "Received JSON response: {}\n\n".format(json_result)
            )
            result = json.loads(json_result.decode('ascii'))
        my_sock.sendall(b'?WATCH={"enable":false}')
        my_sock.close()
        return True, result

    def bist_tpm(self):
        """
        BIST for TPM (Trusted Platform Module)
        Description:!!! I don't know how this will be tested !!!
        External Equipment: None
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'idk': Returns some stuff??
            }
        """
        assert 'tpm' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_clock_int(self):
        """
        BIST for clock lock from internal source
        Description: Checks to see if the N3xx can lock to the internal reference clock
        External Equipment: None
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'locked': Return TRUE if able to lock to internal reference, else FALSE
            }
        """
        assert 'clock_int' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_clock_ext(self):
        """
        BIST for clock lock from external source
        Description: Checks to see if the N3xx can lock to the external reference clock
        External Equipment: Reference Source outputted into the "Ref In"
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'locked': Return TRUE if able to lock to external reference, else FALSE
            }
        """
        assert 'clock_ext' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_usbhost(self):
        """
        BIST for USB host functionality
        Description:!!! I'm considering removing this test !!!
        External Equipment: USB device attached to USB hOST
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'devices': [] Returns all devices seen on USB hub
            }
        """
        assert 'usbhost' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_sfp(self):
        """
        BIST for SFP+ ports:
        Description: Uses one SFP+ port to test the other. Returns the speeds
        External Equipment: External loopback cable between SFP+ port 0 and 1 is required
        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'speed01': Returns speed in Gbit/s from port 0 to port 1
            'speed10': Returns speed in Gbit/s from port 1 to port 0
            }
        """
        assert 'sfp' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_gpio(self):
        """
        BIST for GPIO
        Description: Writes and reads the values to the GPIO
        Needed Equipment: External loopback as follows
            GPIO
            0<->6
            1<->7
            2<->8
            3<->9
            4<->10
            5<->11

        JSON{
            'status': Return TRUE if no errors occurred, else FALSE
            'expected': Return expected read values
            'returned': Return read values
            }
        """
        assert 'gpio' in self.tests_to_run
        sys.stderr.write("Test not implemented.\n")
        return True, {}

    def bist_temp(self):
        """
        BIST for temperature sensors
        Description: Reads the temperature sensors on the motherboards and
        returns their values in mC

        Return dictionary:
        - <thermal-zone-name>: temp in mC
        """
        assert 'temp' in self.tests_to_run
        if self.args.dry_run:
            return True, {'fpga-thermal-zone': 30000}
        context = pyudev.Context()
        result = {
            device.attributes.get('type').decode('ascii'): \
                    int(device.attributes.get('temp').decode('ascii'))
            for device in context.list_devices(subsystem='thermal')
            if 'temp' in device.attributes.available_attributes \
                    and device.attributes.get('temp') is not None
        }
        return True, result

    def bist_fan(self):
        """
        BIST for temperature sensors
        Description: Reads the RPM values of the fans on the motherboard

        Return dictionary:
        - <fan-name>: Fan speed in RPM

        External Equipment: None
        """
        assert 'fan' in self.tests_to_run
        if self.args.dry_run:
            return True, {'cooling_device0': 10000}
        context = pyudev.Context()
        result = {
            device.sys_name: int(device.attributes.get('cur_state'))
            for device in context.list_devices(subsystem='thermal')
            if 'cur_state' in device.attributes.available_attributes \
                    and device.attributes.get('cur_state') is not None
        }
        return True, result

def main():
    " Go, go, go! "
    return N310BIST().run()

if __name__ == '__main__':
    exit(not main())