import os, json, requests
from getpass import getpass
from APIHandler import APIHandler


class Layer:
    """
    Layer(layer_data) - where layer_data is a Python dictionary.

    Meant to be used for eLabFTW Experiments of the "PLD Deposition" category.

    eLabFTW experiments contain most of the data required by the NeXus file - although every layer is on a different eLab entry;
    unfortunately, some data like the target's chemical formula must be retrieved through additional HTTP requests.
    Attributes 'target_elabid', 'rheed_system_elabid' and 'laser_system_elabid' contain elabid's for these resources, which are all items.
    """

    def __init__(self, layer_data):
        """
        Properties/Attributes:
            Too many to list.
        """
        try:
            self.elabid = layer_data["id"]
            self.operator = layer_data["fullname"]
            self.extra = layer_data["metadata_decoded"]["extra_fields"]
            self.uploads = layer_data["uploads"]  # dict
            self.layer_number = self.extra["Layer Progressive Number"][
                "value"
            ]  # integer
            self.target_elabid = self.extra["Target"]["value"]  # elabid
            self.laser_system_elabid = self.extra["Laser System"]["value"]  # elabid
            self.chamber_elabid = self.extra["Chamber"]["value"]  # elabid
            self.rheed_system_elabid = self.extra["RHEED System"]["value"]  # elabid
            self.deposition_time = self.extra["Duration"]["value"]
            self.deposition_time_unit = self.extra["Duration"]["unit"]
            self.repetition_rate = self.extra["Repetition rate"]["value"]
            self.repetition_rate_unit = self.extra["Repetition rate"]["unit"]
            try:
                self.number_of_pulses = (
                    float(self.deposition_time) * float(self.repetition_rate)
                ).__floor__()
            except ValueError:
                # Since number_of_pulses is required, if it can't be calculated raise error:
                raise ValueError("""
    Fatal: either Duration or Repetition Rate are empty or invalid.
    This has to be an error, since these fields are required by the NeXus standard.
    Please edit your eLabFTW entry and retry.
                """)
            self.temperature = self.extra["Heater temperature"][
                "value"
            ]  # Note: this field used to have a trailing space in its name
            self.temperature_unit = self.extra["Heater temperature"]["unit"]
            self.process_pressure = self.extra["Process pressure"][
                "value"
            ]  # Note: this field used to have a trailing space in its name
            self.process_pressure_unit = self.extra["Process pressure"]["unit"]
            self.heating_method = self.extra["Heating Method"]["value"]
            self.layer_thickness = self.extra["Thickness"]["value"]
            self.layer_thickness_unit = self.extra["Thickness"]["unit"]
            self.buffer_gas = self.extra["Buffer gas"]["value"]
            self.heater_target_distance = self.extra["Heater-target distance"]["value"]
            self.heater_target_distance_unit = self.extra["Heater-target distance"][
                "unit"
            ]
            self.laser_fluence = self.extra["Laser Intensity"][
                "value"
            ]  # here fluence = intensity
            self.laser_fluence_unit = "J/(s cm^2)"
            self.laser_spot_area = self.extra["Spot Area"]["value"]
            self.laser_spot_area_unit = "mm^2"
            try:
                self.laser_energy = (
                    float(self.laser_fluence) * float(self.laser_spot_area) / 100
                ).__round__(3)
                self.laser_energy_unit = "J/s"
            except ValueError:
                # Since laser_energy is NOT required, if it can't be calculated warn user but allow the software to continue execution:
                print("""
    Warning: either Laser Intensity or Spot Area are empty or invalid.
    If you think this is an error, please edit your eLabFTW entry and retry.
    Setting Laser Energy to NoneType.
                """)
                # Placeholder
                self.laser_energy = "N/A"
                self.laser_energy_unit = "J/s"
            # Laser rasternig section
            self.laser_rastering_geometry = self.extra["Laser Rastering Geometry"][
                "value"
            ]
            self.laser_rastering_positions = self.extra["Laser Rastering Position"][
                "value"
            ]
            self.laser_rastering_velocities = self.extra["Laser Rastering Speed"][
                "value"
            ]
            # Pre annealing section
            self.pre_annealing_ambient_gas = self.extra["Buffer gas Pre"]["value"]
            self.pre_annealing_pressure = self.extra["Process pressure Pre"]["value"]
            self.pre_annealing_temperature = self.extra["Heater temperature Pre"][
                "value"
            ]
            self.pre_annealing_duration = self.extra["Duration Pre"]["value"]
            self.pre_annealing_pressure_unit = self.extra["Process pressure Pre"][
                "unit"
            ]
            self.pre_annealing_temperature_unit = self.extra["Heater temperature Pre"][
                "unit"
            ]
            self.pre_annealing_duration_unit = self.extra["Duration Pre"]["unit"]
            # Post annealing section
            self.post_annealing_ambient_gas = self.extra["Buffer gas PA"]["value"]
            self.post_annealing_pressure = self.extra["Process pressure PA"]["value"]
            self.post_annealing_temperature = self.extra["Heater temperature PA"][
                "value"
            ]
            self.post_annealing_duration = self.extra["Duration PA"]["value"]
            self.post_annealing_pressure_unit = self.extra["Process pressure PA"][
                "unit"
            ]
            self.post_annealing_temperature_unit = self.extra["Heater temperature PA"][
                "unit"
            ]
            self.post_annealing_duration_unit = self.extra["Duration PA"]["unit"]

            # Rejected but suggested by the NeXus standard:
            # self.laser_rastering_coefficients = None
        except KeyError as k:
            # Some keys are not required and can be called through the .get() method - which is permissive and allows null values;
            # Other keys are required so if they can't be called (invalid or null) raise error and stop execution of the program:
            raise KeyError(
                f'The provided dictionary lacks a "{k}" key. Check the deposition layer entry on eLabFTW and make sure you used the correct Experiment template.'
            )
        # Optional
        self.start_time = layer_data.get("created_at") or None
        self.description = layer_data.get("body") or None

    def get_instruments(self, api_key, ELABFTW_API_URL):
        """
        Retruns a dictionary of all the instruments used to create the layer.
        The format of the dictionary is:
            {
                "laser_system": str,
                "deposition_chamber": str,
                "rheed_system": str
            }

        Args:
            api_key: str:           A valid API key for the eLabFTW instance where the data is stored, with permissions to access the relevant entries.
                                    eLabFTW's API keys are well documented here: https://doc.elabftw.net/docs/usage/api/.
                                    If you don't have an API key and are uncapable of creating one, contact your eLabFTW administrator.
                                    Or RTFM and create one yourself, it's not that hard.
            ELABFTW_API_URL: str:   URL for the API root endpoint of the eLabFTW instance. Ends with '/api/v2' - no trailing slash.
        """
        raw_lasersys_data = APIHandler(api_key, ELABFTW_API_URL).get_entry_from_elabid(
            self.laser_system_elabid, entryType="items"
        )
        raw_chamber_data = APIHandler(api_key, ELABFTW_API_URL).get_entry_from_elabid(
            self.chamber_elabid, entryType="items"
        )
        raw_rheedsys_data = APIHandler(api_key, ELABFTW_API_URL).get_entry_from_elabid(
            self.rheed_system_elabid, entryType="items"
        )
        instruments_used = {
            "laser_system": raw_lasersys_data.get("title") or None,
            "deposition_chamber": raw_chamber_data.get("title") or None,
            "rheed_system": raw_rheedsys_data.get("title") or None,
        }
        return instruments_used

    def list_attachments(self):
        """
        Returns a dictionary of all the attachments linked to the layer, where:
            * Each key is the attachment's progressive ID (0, 1...);
            * Each value is a dictionary containing the attachment's elabid, filename, hashname and related experiment elabid (= self.elabid).

        Data is already in layer_data, so the API key is unrequired. Same goes for:
            * fetch_textual_uploads() - no arguments;
            * fetch_images() - no arguments.

        Exception: returns {} (empty dictionary) if no uploads/attachments on Layer.
        """
        # Remember: Layers are experiments, so we only need to look for attachments in the experiment endpoint.
        if self.uploads == []:
            return {}
        attachments = {
            self.uploads.index(attachment): {
                "id": attachment["id"],
                "filename": attachment["real_name"],
                "hashname": attachment["long_name"],
                "related_experiment": attachment["item_id"],
            }
            for attachment in self.uploads
        }
        return attachments

    def fetch_textual_uploads(self):
        """
        Starting from the list of attachments, filters out and returns a list of the textual uploads linked to the layer, which can be either plain text, csv, tsv etc.
        Returns only their names, so that the user may select which one to import into the NeXus file as a dataset.

        It only looks for .txt, .csv and .tsv files, although it could be easily modified to include other formats.
        It is also file extension-sensitive, so anything not ending with .txt, .csv or .tsv won't be retrieved.
        That's because the API (v5.3.11) doesn't provide MIME Type or similar metadata on the attachments, so the only way to know if an attachment is an image or not is through its filename.
        """
        attachments = self.list_attachments()
        textual_uploads = {
            attachment: attachments[attachment]
            for attachment in attachments
            if attachments[attachment]["filename"][-4:] in (".txt", ".csv", ".tsv")
        }
        return textual_uploads

    def fetch_images(self):
        """
        Starting from the list of attachments, filters out and returns a Starting from the list of attachments, filters out and returns a list of all the (PNG or BMP) images attached to the layer.
        Hopefully one of them is a RHEED pattern.
        Returns only their names, so that the user may select which one to import into the NeXus file as a RHEED acquisition.

        It only looks for .png and .bmp files, although it could be easily modified to include other formats.
        It is also file extension-sensitive, so anything not ending with .png or .bmp won't be retrieved, even if it's an actual image.
        That's because the API (v5.3.11) doesn't provide MIME Type or similar metadata on the attachments, so the only way to know if an attachment is an image or not is through its filename.
        """
        attachments = self.list_attachments()
        images = {
            attachment: attachments[attachment]
            for attachment in attachments
            if attachments[attachment]["filename"][-4:] in (".png", ".bmp")
        }
        return images


class Entrypoint:
    """
    Entrypoint(sample_data) - where sample_data is a Python dictionary.

    Meant to be used for eLabFTW Resources of the "Sample" category.

    The entrypoint is the starting point of the process of resolving the data chain.
    The entrypoint must be a dictionary containing the data of a sample, created directly from the JSON of the item endpoint on eLabFTW - which can be done through the function get_entry_from_elabid.
    """

    def __init__(self, sample_data):
        """
        Properties/Attributes:
            * name: str:                        Name of the sample. Fairly important, and always present unless someone screws up REALLY bad.
            * linked_items: dict:               Dictionary generated by eLabFTW containing metadata on the items linked to the entrypoint.
            * batch_elabid: int:                eLabFTW internal id of the batch of the substrate used as the foundation of the sample.
            * proposal: int:                    eLabFTW internal id of the proposal linked to the sample.
            * linked_experiments: dict:         Dictionary generated by eLabFTW containing metadata on the experiments linked to the entrypoint.
            * linked_experiments_elabid: list:  List of eLabFTW internal id's of the experiments linked to the entrypoint.
        """
        try:
            self.name = sample_data["title"]
            self.extra = sample_data["metadata_decoded"]["extra_fields"]
            self.linked_items = sample_data["items_links"]  # dict
            self.batch_elabid = self.extra["Substrate batch"]["value"]  # elabid
            self.proposal = self.extra["Proposal"].get("value") or None  # proposal
            self.linked_experiments = sample_data["related_experiments_links"]  # dict
            self.linked_experiments_elabid = [
                i["entityid"] for i in self.linked_experiments
            ]  # list of elabid
        except KeyError as k:
            # Some keys are not required and can be called through the .get() method - which is permissive and allows null values;
            # Other keys are required so if they can't be called (invalid or null) raise error and stop execution of the program:
            raise KeyError(
                f'The provided dictionary lacks a "{k}" key. Check the sample entry on eLabFTW and make sure you used the correct Resource template.'
            )


class Material:
    """
    Material(material_data) - where material_data is a Python dictionary.

    Meant to be used for eLabFTW Resources of either the "PLD Target" or the "Substrate" categories.

    Both a PLD Target and a Substrate are materials made of a certain compound, of which we want to know:
        * Name and formula;
        * Shape and dimensions;
        * Misc.
    """

    def __init__(self, material_data):
        """
        Properties/Attributes:
            * name: str:            Name of the material.
            * compound_elabid: int: eLabFTW internal id of the compound.
            * dimensions: str:      Dimensions of the material, in standard format.
                                    The class recognizes the unit of measurement and acts consequently.
            * dimensions_unit: str: Unit of measurement - either "mm x mm", "inches" or None.
        """
        try:
            self.name = material_data["title"]  # required
            self.extra = material_data["metadata_decoded"]["extra_fields"]
            self.compound_elabid = self.extra["Compound"]["value"]
            self.dimensions = str(
                self.extra["Size"]["value"]
            )  # strings have a .count() method
            if self.dimensions.count("mm") == 2:
                self.dimensions_unit = "mm x mm"
            elif self.dimensions[-1] == '"':
                self.dimensions_unit = "inches"
            else:
                self.dimensions_unit = None
        except KeyError as k:
            # Some keys are not required and can be called through the .get() method - which is permissive and allows null values;
            # Other keys are required so if they can't be called (invalid or null) raise error and stop execution of the program:
            raise KeyError(
                f'The provided dictionary lacks a "{k}" key. Check the target/substrate entry on eLabFTW and make sure you used the correct Resource template.'
            )

    def get_compound_data(self, apikey, ELABFTW_API_URL):
        """
        Returns a dictionary with the relevant data on the compound of which the material is made.
        The format of the dictionary is:
            {
                "name": str,
                "chemical_formula": str,
                "cas_number": str
            }

        Args:
            api_key: str:           A valid API key for the eLabFTW instance where the data is stored, with permissions to access the relevant entries.
                                    eLabFTW's API keys are well documented here: https://doc.elabftw.net/docs/usage/api/.
                                    If you don't have an API key and are uncapable of creating one, contact your eLabFTW administrator.
                                    Or RTFM and create one yourself, it's not that hard.
            ELABFTW_API_URL: str:   URL for the API root endpoint of the eLabFTW instance. Ends with '/api/v2' - no trailing slash.
        """
        raw_compound_data = APIHandler(apikey, ELABFTW_API_URL).get_entry_from_elabid(
            self.compound_elabid, entryType="items"
        )
        name = raw_compound_data["title"]
        extra = raw_compound_data["metadata_decoded"]["extra_fields"]
        formula = extra.get("Chemical formula")
        cas = extra.get("CAS number ") or {"value": None}
        compound_data = {
            "name": name,
            "chemical_formula": formula.get("value"),
            "cas_number": cas.get("value"),
        }
        return compound_data

    def get_compound_formula(self, apikey, ELABFTW_API_URL):
        """
        Returns a string with the chemical formula of the compound.

        Args:
            api_key: str:           A valid API key for the eLabFTW instance where the data is stored, with permissions to access the relevant entries.
                                    eLabFTW's API keys are well documented here: https://doc.elabftw.net/docs/usage/api/.
                                    If you don't have an API key and are uncapable of creating one, contact your eLabFTW administrator.
                                    Or RTFM and create one yourself, it's not that hard.
            ELABFTW_API_URL: str:   URL for the API root endpoint of the eLabFTW instance. Ends with '/api/v2' - no trailing slash.
        """
        formula = self.get_compound_data(apikey, ELABFTW_API_URL).get(
            "chemical_formula"
        )
        return formula


class Substrate(Material):
    """
    Substrate(material_data) - where material_data is a Python dictionary.

    Inherits from Material and it's meant to be used exclusively for eLabFTW Resources of the "Substrate" category.
    """

    def __init__(self, material_data):
        """
        Properties/Attributes common to all Materials:
            * name: str:            Name of the material.
            * compound_elabid: int: eLabFTW internal id of the compound.
            * dimensions: str:      Dimensions of the material, in standard format.
                                    The class recognizes the unit of measurement and acts consequently.
            * dimensions_unit: str: Unit of measurement - either "mm x mm", "inches" or None.

        Specific properties/attributes:
            * orientation: str:
            * miscut_angle: str:
            * miscut_angle_unit: str:
            * miscut_direction: str:
            * thickness: str:
            * thickness_unit: str:
            * surface_treatment: str:
            * manufacturer: str:
            * batch_id: str:
        """
        super().__init__(material_data)
        try:
            self.orientation = self.extra["Orientation"]["value"]
            self.miscut_angle = self.extra["Miscut Angle"]["value"]
            self.miscut_angle_unit = self.extra["Miscut Angle"]["unit"]
            self.miscut_direction = self.extra["Miscut Direction"]["value"]
            # Not present (yet) on eLabFTW for Substrates:
            self.thickness = ""  # self.extra["Thickness"]["value"]
            self.thickness_unit = "μm"  # self.extra["Thickness"]["unit"]
            self.surface_treatment = self.extra["Surface treatment"]["value"]
            self.manufacturer = self.extra["Supplier"]["value"]
            self.batch_id = self.extra["Batch ID"]["value"]
        except KeyError as k:
            raise KeyError(
                f'The provided dictionary lacks a "{k}" key - which is specific for substrates. Check the {self.name} substrate entry on eLabFTW and make sure you used the correct Resource template.'
            )


class Target(Material):
    """
    Target(material_data) - where material_data is a Python dictionary.

    Inherits from Material and it's meant to be used exclusively for eLabFTW Resources of the "PLD Target" category.
    """

    def __init__(self, material_data):
        """
        Properties/Attributes common to all Materials:
            * name: str:            Name of the material.
            * compound_elabid: int: eLabFTW internal id of the compound.
            * dimensions: str:      Dimensions of the material, in standard format.
                                    The class recognizes the unit of measurement and acts consequently.
            * dimensions_unit: str: Unit of measurement - either "mm x mm", "inches" or None.

        Specific properties/attributes:
            * thickness: str:
            * thickness_unit: str:
            * shape: str:
            * solid_form: str:
            * manufacturer: str:
        """
        super().__init__(material_data)
        try:
            self.thickness = self.extra["Thickness"]["value"]
            self.thickness_unit = self.extra["Thickness"]["unit"]
            self.shape = self.extra["shape"]["value"]
            self.solid_form = self.extra["Solid form"]["value"]
            self.manufacturer = self.extra["Supplier"]["value"]
        except KeyError as k:
            raise KeyError(
                f'The provided dictionary lacks a "{k}" key - which is specific for PLD targets. Check the {self.name} target entry on eLabFTW and make sure you used the correct Resource template.'
            )
        # Non-required attributes:
        self.description = material_data.get("body") or ""


class Proposal:
    """
    Proposal(proposal_data) - where proposal_data is a Python dictionary.

    Recovers only the relevant info on a proposal linked to the entrypoint sample.
    Which currently is just its name.

    If the name starts with "Proposal " (space included) that gets omitted from the output.
    """

    def __init__(self, proposal_data):
        """
        Properties/Attributes:
            * name: str:    Name of the proposal.
                            If the name starts with "Proposal " (space included) that gets omitted from the output.
        """
        if "Proposal " in proposal_data["title"]:
            self.name = proposal_data["title"].replace("Proposal ", "")
        else:
            self.name = proposal_data["title"]


if __name__ == "__main__":
    # head = APIHandler("MyApiKey-123456789abcdef")
    # print(f"Example header:\n\t{head.header}\n")
    # print("Warning: you're not supposed to be running this as the main program.")
    api_key = getpass("Paste API key here [no echo]: ")
    ELABFTW_API_URL = input("Enter a valid eLabFTW API URL (ends with '/api/v2)': ")
    handler = APIHandler(api_key, ELABFTW_API_URL)
    exp58 = handler.get_entry_from_elabid(elabid=58, entryType="experiments")
    layer58 = Layer(exp58)
    print(layer58.list_attachments())
    print(layer58.fetch_textual_uploads())
    print(layer58.fetch_images())
