import os
import sys

from mp_api.client import MPRester
from pymatgen.core.structure import Structure


formula = None
output_format = "poscar"


def usage():
    print()
    print(f"Usage: python {sys.argv[0]} formula [output_format]\n")

def update_vars():
    global formula, output_formuat

    argv = sys.argv
    narg = len(argv)
    if narg <= 1:
        print("\nError: Chemical formula must be given as the first arg\n")
        input("Press ENTER to terminate>>\n")
        exit()

    formula = argv[1]
    if narg >= 3:
        output_format = argv[2]


def get_structure(formula):
    API_KEY = os.getenv('MP_APIKEY')
    if API_KEY is None:
        print("\nError: Can not get MP API Key from the environment var MP_APIKEY\n")
        return None


    mpr = MPRester(API_KEY)
    if mpr is None:
        print(f"\nError: Can not get MPRester using the given API_KEY [{API_KEY}]\n")
        input("Press ENTER to terminate>>\n")
        exit()

    search_results = mpr.materials.search(formula = formula)
    if not search_results:
        print(f"No data found for {formula}")
        return None

    structures = []
    for res in search_results:
        material_id = res.material_id
        print(f"Found material_id for {formula}: {material_id}")

        material_data = mpr.materials.search(material_ids=[material_id])
        if not material_data:
            print(f"No structure data found for material_id {material_id}")
            continue

        structure_dict = material_data[0].structure.as_dict()
        structure = Structure.from_dict(structure_dict)
        structures.append(structure)

    return structures


def save_structures(structures, output_format, formula):
    for i, structure in enumerate(structures):
        if output_format == "poscar":
            if i == 0:
                output_path = f"POSCAR"
            else:
                output_path = f"POSCAR_{i+1}"
        else:
            if i == 0:
                output_path = f"{formula}.{output_format}"
            else:
                output_path = f"{formula}_{i+1}.{output_format}"

        print(f"Saving {formula} structure to [{output_path}]...")
        structure.to(filename = output_path, fmt = output_format)
        

def main():
    update_vars()

    structures = get_structure(formula)
    if structures:
        save_structures(structures, output_format, formula)

    usage()
    print()
    input("Press ENTER to terminate>>\n")
    exit()


if __name__ == "__main__":
    main()