# !/usr/bin/env python3

import os, warnings, random, shutil
from pathlib import Path
from dotenv import load_dotenv
from ase.io import read, write
from ase.build import surface
from mp_api.client import MPRester
from pymatgen.core import Structure
from pymatgen.io.vasp import Poscar, Kpoints
from pymatgen.analysis.defects import generators
from pymatgen.io.vasp.sets import (
    MPStaticSet, 
    MPRelaxSet, 
    MPNonSCFSet, 
    MPScanRelaxSet, 
    MPScanStaticSet,
    MPMDSet,
    NEBSet,
    MVLElasticSet,
    )

from masgent import schemas
from masgent.utils import (
    write_comments,
    color_print,
    ask_for_mp_api_key,
    validate_mp_api_key,
    )

# Do not show warnings
warnings.filterwarnings('ignore')

# Track whether Materials Project key has been checked during this process
_mp_key_checked = False

def with_metadata(input: schemas.ToolMetadata):
    '''
    Decorator to add metadata to tool functions.
    '''
    def decorator(func):
        func._tool_metadata = input
        return func
    return decorator

@with_metadata(schemas.ToolMetadata(
    name='List Files',
    description='List all files in the current session runs directory.',
    requires=[],
    optional=[],
    defaults={},
    prereqs=[],
))
def list_files() -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

    file_list = []
    base_dir = Path(runs_dir)
    for item in base_dir.rglob('*'):
        if item.is_file():
            file_list.append(str(item.relative_to(base_dir)))
    return {
        'status': 'success',
        'message': f'Found {len(file_list)} files in the current session runs directory.',
        'files': file_list,
    }

@with_metadata(schemas.ToolMetadata(
    name='Read File',
    description='Read a file from the current session runs directory.',
    requires=['name'],
    optional=[],
    defaults={},
    prereqs=[],
))
def read_file(name: str) -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
    base_dir = Path(runs_dir)

    try:
        with open(base_dir / name, "r") as f:
            content = f.read()
        return {
            'status': 'success',
            'message': f'File {name} read successfully.',
            'content': content,
        }
    except Exception as e:
        return {
            'status': 'error',
            'message': f'An error occurred while reading file {name}: {e}',
        }

@with_metadata(schemas.ToolMetadata(
    name='Rename File',
    description='Rename a file in the current session runs directory.',
    requires=['name', 'new_name'],
    optional=[],
    defaults={},
    prereqs=[],
))
def rename_file(name: str, new_name: str) -> dict:
    runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
    base_dir = Path(runs_dir)

    try:
        new_path = base_dir / new_name
        if not str(new_path).startswith(str(base_dir)):
            return {
                'status': 'error',
                'message': f'Renaming to {new_name} would move the file outside the session runs directory, which is not allowed.',
            }

        os.makedirs(new_path.parent, exist_ok=True)
        shutil.copy2(base_dir / name, new_path)
        return {
            'status': 'success',
            'message': f'File {name} renamed to {new_name} successfully.',
        }
    except Exception as e:
        return {
            'status': 'error',
            'message': f'An error occurred while renaming file {name} to {new_name}: {e}',
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate POSCAR',
    description='Generate POSCAR file from Materials Project database based on chemical formula.',
    requires=['formula'],
    optional=[],
    defaults={},
    prereqs=[],
))
def generate_vasp_poscar(input: schemas.GenerateVaspPoscarSchema) -> dict:
    '''
    Generate VASP POSCAR file from Materials Project database.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar with input: {input}', 'green')
    
    formula_list = input.formula_list

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        poscar_paths = []
        
        for formula in formula_list:
            
            poscar_paths.append(os.path.join(runs_dir, 'POSCAR'))

            # Ensure Materials Project API key exists and validate it only once per process
            load_dotenv(dotenv_path='.env')

            global _mp_key_checked
            if not _mp_key_checked:
                if 'MP_API_KEY' not in os.environ:
                    ask_for_mp_api_key()
                else:
                    # color_print('[Info] Materials Project API key found in environment.\n', 'green')
                    validate_mp_api_key(os.environ['MP_API_KEY'])
                _mp_key_checked = True
            
            with MPRester(mute_progress_bars=True) as mpr:
                docs = mpr.materials.summary.search(formula=formula)
                if not docs:
                    return {
                        'status': 'error',
                        'message': f'No materials found in Materials Project database for formula: {formula}'
                    }
                
                mid = docs[0].material_id   # pick the first match
                structure = mpr.get_structure_by_material_id(mid)
                poscar = Poscar(structure)

            # Also write to the main target directory for easy access
            poscar.write_file(os.path.join(runs_dir, 'POSCAR'), direct=True)

            comments = f'# Generated by Masgent from Materials Project entry {mid}.'
            write_comments(os.path.join(runs_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated POSCAR(s) in {poscar_paths}.',
            'poscar_paths': poscar_paths,
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'POSCAR generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate VASP Inputs (INCAR, KPOINTS, POTCAR, POSCAR)',
    description='Generate VASP input files (INCAR, KPOINTS, POTCAR, POSCAR) from a given POSCAR file using pymatgen input sets (MPRelaxSet, MPStaticSet, MPNonSCFSet, MPScanRelaxSet, MPScanStaticSet, MPMDSet, NEBSet, MVLElasticSet).',
    requires=['vasp_input_sets'],
    optional=['poscar_path', 'only_incar'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'only_incar': False,
        },
    prereqs=[],
))
def generate_vasp_inputs_from_poscar(input: schemas.GenerateVaspInputsFromPoscar) -> dict:
    '''
    Generate VASP input files (INCAR, KPOINTS, POTCAR, POSCAR) from a given POSCAR file using pymatgen input sets.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_inputs_from_poscar with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    vasp_input_sets = input.vasp_input_sets
    only_incar = input.only_incar

    VIS_MAP = {
        'MPRelaxSet': MPRelaxSet,
        'MPStaticSet': MPStaticSet,
        'MPNonSCFSet': MPNonSCFSet,
        'MPScanRelaxSet': MPScanRelaxSet,
        'MPScanStaticSet': MPScanStaticSet,
        'MPMDSet': MPMDSet,
        'NEBSet': NEBSet,
        'MVLElasticSet': MVLElasticSet,
    }
    vis_class = VIS_MAP[vasp_input_sets]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        vasp_inputs_dir = os.path.join(runs_dir, f'vasp_inputs/{vasp_input_sets}')
        os.makedirs(vasp_inputs_dir, exist_ok=True)

        structure = Structure.from_file(poscar_path)
        vis = vis_class(structure)

        if only_incar:
            vis.incar.write_file(os.path.join(vasp_inputs_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(vasp_inputs_dir, 'POSCAR'))
            incar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
            write_comments(os.path.join(vasp_inputs_dir, 'INCAR'), 'incar', incar_comments)
            return {
                'status': 'success',
                'message': f'Generated INCAR based on {vasp_input_sets} in {vasp_inputs_dir}.',
                'incar_path': os.path.join(vasp_inputs_dir, 'INCAR'),
            }
        
        vis.incar.write_file(os.path.join(vasp_inputs_dir, 'INCAR'))
        vis.poscar.write_file(os.path.join(vasp_inputs_dir, 'POSCAR'))
        vis.kpoints.write_file(os.path.join(vasp_inputs_dir, 'KPOINTS'))
        vis.potcar.write_file(os.path.join(vasp_inputs_dir, 'POTCAR'))

        incar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'INCAR'), 'incar', incar_comments)
        poscar_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'POSCAR'), 'poscar', poscar_comments)
        kpoints_comments = f'# Generated by Masgent using {vasp_input_sets} set provided by Materials Project.'
        write_comments(os.path.join(vasp_inputs_dir, 'KPOINTS'), 'kpoints', kpoints_comments)
        
        return {
            'status': 'success',
            'message': f'Generated VASP input files based on {vasp_input_sets} in {vasp_inputs_dir}.',
            'incar_path': os.path.join(vasp_inputs_dir, 'INCAR'),
            'kpoints_path': os.path.join(vasp_inputs_dir, 'KPOINTS'),
            'potcar_path': os.path.join(vasp_inputs_dir, 'POTCAR'),
            'poscar_path': os.path.join(vasp_inputs_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP input files generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Convert Structure Format',
    description='Convert structure files between different formats (CIF, POSCAR, XYZ).',
    requires=['input_path', 'input_format', 'output_format'],
    optional=[],
    defaults={},
    prereqs=[],
))
def convert_structure_format(input: schemas.ConvertStructureFormatSchema) -> dict:
    '''
    Convert structure files between different formats (CIF, POSCAR, XYZ).
    '''
    # color_print(f'\n[Debug: Function Calling] convert_structure_format with input: {input}', 'green')
    
    input_path = input.input_path
    input_format = input.input_format
    output_format = input.output_format
    
    format_map = {
        "POSCAR": "vasp",
        "CIF": "cif",
        "XYZ": "xyz"
    }

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        convert_dir = os.path.join(runs_dir, 'convert')
        os.makedirs(convert_dir, exist_ok=True)

        atoms = read(input_path, format=format_map[input_format])
        filename_wo_ext = os.path.splitext(os.path.basename(input_path))[0]
        # Ignore the POSCAR, do not add extension
        if output_format == 'POSCAR':
            output_path = os.path.join(convert_dir, 'POSCAR')
        else:
            output_path = os.path.join(convert_dir, f'{filename_wo_ext}.{output_format.lower()}')
        write(output_path, atoms, format=format_map[output_format])

        return {
            'status': 'success',
            'message': f'Converted structure saved to {output_path}.',
            'output_path': output_path
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Structure conversion failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Convert POSCAR Coordinates',
    description='Convert POSCAR between direct and cartesian coordinates.',
    requires=['to_cartesian'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def convert_poscar_coordinates(input: schemas.ConvertPoscarCoordinatesSchema) -> dict:
    '''
    Convert POSCAR between direct and cartesian coordinates.
    '''
    # color_print(f'\n[Debug: Function Calling] convert_poscar_coordinates with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    to_cartesian = input.to_cartesian

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        convert_dir = os.path.join(runs_dir, 'convert')
        os.makedirs(convert_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path)
        poscar = Poscar(structure)
        poscar.write_file(os.path.join(convert_dir, 'POSCAR'), direct=not to_cartesian)
        coord_type = 'Cartesian' if to_cartesian else 'Direct'
        comments = f'# Generated by Masgent converted to {coord_type} coordinates.'
        write_comments(os.path.join(convert_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Converted POSCAR to {coord_type} coordinates in {convert_dir}.',
            'poscar_path': os.path.join(convert_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'POSCAR coordinate conversion failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Customize KPOINTS',
    description='Customize VASP KPOINTS from POSCAR with specified accuracy level (Low, Medium, High).',
    requires=['accuracy_level'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def customize_vasp_kpoints_with_accuracy(input: schemas.CustomizeVaspKpointsWithAccuracy) -> dict:
    '''
    Customize VASP KPOINTS from POSCAR with specified accuracy level.
    '''
    # color_print(f'\n[Debug: Function Calling] customize_vasp_kpoints_with_accuracy with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    accuracy_level = input.accuracy_level
    
    DENSITY_MAP = {
        'Low': 1000,
        'Medium': 3000,
        'High': 5000,
    }
    kppa = DENSITY_MAP[accuracy_level]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        structure = Structure.from_file(poscar_path)
        kpoints = Kpoints.automatic_density(structure, kppa=kppa)
        kpoints.write_file(os.path.join(runs_dir, 'KPOINTS'))
        comments = f'# Generated by Masgent with {accuracy_level} accuracy (Grid Density = {kppa} / number of atoms)'
        write_comments(os.path.join(runs_dir, 'KPOINTS'), 'kpoints', comments)
        
        return {
            'status': 'success',
            'message': f'Updated KPOINTS with {accuracy_level} accuracy in {runs_dir}.',
            'kpoints_path': os.path.join(runs_dir, 'KPOINTS'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP KPOINTS generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Vacancy Defects',
    description='Generate VASP POSCAR with vacancy defects.',
    requires=['original_element', 'defect_amount'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_vacancy_defects(input: schemas.GenerateVaspPoscarWithVacancyDefects) -> dict:
    '''
    Generate VASP POSCAR with vacancy defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_vacancy_defects with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    original_element = input.original_element
    defect_amount = input.defect_amount

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        all_indices = [i for i, atom in enumerate(atoms) if atom.symbol == original_element]
        if isinstance(defect_amount, float):
            num_defects = max(1, int(defect_amount * len(all_indices)))
        elif isinstance(defect_amount, int):
            num_defects = defect_amount

        vacancy_indices = random.sample(all_indices, num_defects)
        del atoms[vacancy_indices]

        write(os.path.join(defect_dir, 'POSCAR'), atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent with vacancy defects of element {original_element} by randomly removing {num_defects} atoms, be careful to verify structure.'
        write_comments(os.path.join(defect_dir, 'POSCAR'), 'poscar', comments)

        # return f'\nGenerated POSCAR with vacancy defects in {os.path.join(target_dir, "POSCAR")}.'
        return {
            'status': 'success',
            'message': f'Generated POSCAR with vacancy defects in {os.path.join(defect_dir, "POSCAR")}.',
            'poscar_path': os.path.join(defect_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Substitution Defects',
    description='Generate VASP POSCAR with substitution defects.',
    requires=['original_element', 'defect_element', 'defect_amount'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_substitution_defects(input: schemas.GenerateVaspPoscarWithSubstitutionDefects) -> dict:
    '''
    Generate VASP POSCAR with substitution defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_substitution_defects with input: {input}', 'green')

    poscar_path = input.poscar_path
    original_element = input.original_element
    defect_element = input.defect_element
    defect_amount = input.defect_amount

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        all_indices = [i for i, atom in enumerate(atoms) if atom.symbol == original_element]
        if isinstance(defect_amount, float):
            num_defects = max(1, int(defect_amount * len(all_indices)))
        elif isinstance(defect_amount, int):
            num_defects = defect_amount

        substitution_indices = random.sample(all_indices, num_defects)
        for i in substitution_indices:
            atoms[i].symbol = defect_element
        
        write(os.path.join(defect_dir, 'POSCAR'), atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent with substitution defect of element {original_element} to {defect_element} by randomly substituting {num_defects} atoms, be careful to verify structure.'
        write_comments(os.path.join(defect_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated POSCAR with substitution defects in {os.path.join(defect_dir, "POSCAR")}.',
            'poscar_path': os.path.join(defect_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Interstitial (Voronoi) Defects',
    description='Generate VASP POSCAR with interstitial (Voronoi) defects.',
    requires=['defect_element'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_vasp_poscar_with_interstitial_defects(input: schemas.GenerateVaspPoscarWithInterstitialDefects) -> dict:
    '''
    Generate VASP POSCAR with interstitial (Voronoi) defects.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_with_interstitial_defects with input: {input}', 'green')

    poscar_path = input.poscar_path
    defect_element = input.defect_element

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        defect_dir = os.path.join(runs_dir, 'defects')
        os.makedirs(defect_dir, exist_ok=True)
        
        atoms = read(poscar_path, format='vasp')

        # Read atoms from ASE and convert to Pymatgen Structure
        structure = Structure.from_ase_atoms(atoms)
        interstitial_generator = generators.VoronoiInterstitialGenerator().generate(structure=structure, insert_species=[defect_element])
        defect_sites, defect_structures = [], []
        for defect in interstitial_generator:
            defect_sites.append(defect.site.frac_coords)
            defect_structures.append(defect.defect_structure)

        if len(defect_structures) == 0:
            return {
                'status': 'error',
                'message': f'No interstitial sites found for element {defect_element}.',
            }
        else:
            for i, defect_structure in enumerate(defect_structures):
                # Convert back to ASE Atoms for writing
                defect_atoms = defect_structure.to_ase_atoms()
                write(os.path.join(defect_dir, f'POSCAR_{i}'), defect_atoms, format='vasp', direct=True, sort=True)
                comments = f'# Generated by Masgent with interstitial (Voronoi) defect of element {defect_element} at fract. coords {defect_sites[i]}, be careful to verify structure.'
                write_comments(os.path.join(defect_dir, f'POSCAR_{i}'), 'poscar', comments)

        # return f'\nGenerated POSCAR with interstitial (Voronoi) defects in {target_dir}.'
        return {
            'status': 'success',
            'message': f'Generated POSCAR(s) with interstitial (Voronoi) defects in {defect_dir}.',
            'poscar_paths': [os.path.join(defect_dir, f'POSCAR_{i}') for i in range(len(defect_structures))],
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP POSCAR defect generation failed: {str(e)}'
        }

@with_metadata(schemas.ToolMetadata(
    name='Generate Supercell',
    description='Generate supercell from POSCAR based on user-defined 3x3 scaling matrix.',
    requires=['scaling_matrix'],
    optional=['poscar_path'],
    defaults={'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR'},
    prereqs=[],
))
def generate_supercell_from_poscar(input: schemas.GenerateSupercellFromPoscar) -> dict:
    '''
    Generate supercell from POSCAR based on user-defined 3x3 scaling matrix.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_supercell_from_poscar with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    scaling_matrix = input.scaling_matrix

    scaling_matrix_ = [
        [int(num) for num in line.strip().split()] 
        for line in scaling_matrix.split(';')
        ]

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')

        supercell_dir = os.path.join(runs_dir, 'supercell')
        os.makedirs(supercell_dir, exist_ok=True)
        
        structure = Structure.from_file(poscar_path).copy()
        supercell_structure = structure.make_supercell(scaling_matrix_)
        supercell_poscar = Poscar(supercell_structure)
        supercell_poscar.write_file(os.path.join(supercell_dir, 'POSCAR'), direct=True)
        
        comments = f'# Generated by Masgent as supercell with scaling matrix {scaling_matrix}.'
        write_comments(os.path.join(supercell_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated supercell POSCAR in {os.path.join(supercell_dir, "POSCAR")}.',
            'poscar_path': os.path.join(supercell_dir, 'POSCAR'),
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Supercell generation failed: {str(e)}'
        }
    
@with_metadata(schemas.ToolMetadata(
    name='Generate HPC Slurm Script',
    description='Generate HPC Slurm job submission script for VASP calculations.',
    requires=[],
    optional=['partition', 'nodes', 'ntasks', 'walltime', 'jobname', 'mail_type', 'mail_user', 'command'],
    defaults={
        'partition': 'normal',
        'nodes': 1,
        'ntasks': 8,
        'walltime': '01:00:00',
        'jobname': 'masgent_job',
        'command': 'srun vasp_std > vasp.out'
        },
    prereqs=[],
    ))
def generate_vasp_inputs_hpc_slurm_script(input: schemas.GenerateVaspInputsHpcSlurmScript) -> dict:
    '''
    Generate HPC Slurm job submission script for VASP calculations.
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_inputs_hpc_slurm_script with input: {input}', 'green')
    
    partition = input.partition
    nodes = input.nodes
    ntasks = input.ntasks
    walltime = input.walltime
    jobname = input.jobname
    command = input.command

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        scripts = f'''#!/bin/bash
#SBATCH --partition={partition}
#SBATCH --nodes={nodes}
#SBATCH --ntasks={ntasks}
#SBATCH --time={walltime}
#SBATCH --job-name={jobname}
#SBATCH --output={jobname}.out
#SBATCH --error={jobname}.err

# This Slurm script was generated by Masgent, customize as needed.
{command}
'''
        script_path = os.path.join(runs_dir, 'masgent_submit.sh')
        with open(script_path, 'w') as f:
            f.write(scripts)

        return {
            'status': 'success',
            'message': f'Generated Slurm script in {script_path}.',
            'script_path': script_path,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'Slurm script generation failed: {str(e)}'
        }
    
@with_metadata(schemas.ToolMetadata(
    name='Generate VASP POSCAR for suface slab',
    description='Generate VASP POSCAR for surface slab from bulk POSCAR based on Miller indices, vacuum thickness, and slab layers',
    requires=['miller_indices'],
    optional=['poscar_path', 'vacuum_thickness', 'slab_layers'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'vacuum_thickness': 15.0,
        'slab_layers': 4,
        },
    prereqs=[],
))
def generate_vasp_poscar_for_surface_slab(input: schemas.GenerateVaspPoscarForSurfaceSlab) -> dict:
    '''
    Generate VASP POSCAR for surface slab from bulk POSCAR based on Miller indices, vacuum thickness, and slab layers
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_poscar_for_surface_slab with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    miller_indices = input.miller_indices
    vacuum_thickness = input.vacuum_thickness
    slab_layers = input.slab_layers

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        surface_slab_dir = os.path.join(runs_dir, 'surface_slab')
        os.makedirs(surface_slab_dir, exist_ok=True)
        
        bulk_atoms = read(poscar_path, format='vasp')
        slab_atoms = surface(lattice=bulk_atoms, indices=miller_indices, layers=slab_layers, vacuum=vacuum_thickness, tol=1e-10, periodic=True)
        write(os.path.join(surface_slab_dir, 'POSCAR'), slab_atoms, format='vasp', direct=True, sort=True)
        comments = f'# Generated by Masgent as surface slab with Miller indices {miller_indices}, vacuum thickness {vacuum_thickness} Å, and slab layers {slab_layers}.'
        write_comments(os.path.join(surface_slab_dir, 'POSCAR'), 'poscar', comments)

        return {
            'status': 'success',
            'message': f'Generated surface slab POSCAR in {os.path.join(surface_slab_dir, "POSCAR")}.',
            'poscar_path': os.path.join(surface_slab_dir, 'POSCAR'),
        }

    except Exception as e:
        return {
            'status': 'error',
            'message': f'Surface slab POSCAR generation failed: {str(e)}'
        }
    
@with_metadata(schemas.ToolMetadata(
    name='Generate VASP workflow of convergence tests',
    description='Generate VASP workflow of convergence tests for k-points and energy cutoff based on given POSCAR',
    requires=[],
    optional=['poscar_path', 'kpoint_levels', 'encut_levels'],
    defaults={
        'poscar_path': f'{os.environ.get("MASGENT_SESSION_RUNS_DIR")}/POSCAR',
        'kpoint_levels': [1000, 2000, 3000, 4000, 5000],
        'encut_levels': [300, 400, 500, 600, 700],
        },
    prereqs=[],
))
def generate_vasp_workflow_of_convergence_tests(input: schemas.GenerateVaspWorkflowOfConvergenceTests) -> dict:
    '''
    Generate VASP workflow of convergence tests for k-points and energy cutoff based on given POSCAR
    '''
    # color_print(f'\n[Debug: Function Calling] generate_vasp_workflow_of_convergence_tests with input: {input}', 'green')
    
    poscar_path = input.poscar_path
    test_type = input.test_type
    kpoint_levels = input.kpoint_levels
    encut_levels = input.encut_levels

    def test_kpoints():
        for kppa in kpoint_levels:
            vis = MPStaticSet(structure)
            os.makedirs(kpoint_tests_dir, exist_ok=True)
            test_dir = os.path.join(kpoint_tests_dir, f'kppa_{kppa}')
            os.makedirs(test_dir, exist_ok=True)
            vis.incar.write_file(os.path.join(test_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(test_dir, 'POSCAR'))
            vis.potcar.write_file(os.path.join(test_dir, 'POTCAR'))
            incar_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'INCAR'), 'incar', incar_comments)
            poscar_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'POSCAR'), 'poscar', poscar_comments)
            kpoints = Kpoints.automatic_density(structure, kppa=kppa)
            kpoints.write_file(os.path.join(test_dir, 'KPOINTS'))
            kpoint_comments = f'# Generated by Masgent for k-point convergence test with kppa = {kppa}.'
            write_comments(os.path.join(test_dir, 'KPOINTS'), 'kpoints', kpoint_comments)

    def test_encut():
        for encut in encut_levels:
            vis = MPStaticSet(structure, user_incar_settings={'ENCUT': encut})
            os.makedirs(encut_tests_dir, exist_ok=True)
            test_dir = os.path.join(encut_tests_dir, f'encut_{encut}')
            os.makedirs(test_dir, exist_ok=True)
            vis.incar.write_file(os.path.join(test_dir, 'INCAR'))
            vis.poscar.write_file(os.path.join(test_dir, 'POSCAR'))
            vis.potcar.write_file(os.path.join(test_dir, 'POTCAR'))
            vis.kpoints.write_file(os.path.join(test_dir, 'KPOINTS'))
            incar_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'INCAR'), 'incar', incar_comments)
            poscar_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'POSCAR'), 'poscar', poscar_comments)
            kpoint_comments = f'# Generated by Masgent for energy cutoff convergence test with ENCUT = {encut}.'
            write_comments(os.path.join(test_dir, 'KPOINTS'), 'kpoints', kpoint_comments)

    try:
        runs_dir = os.environ.get('MASGENT_SESSION_RUNS_DIR')
        
        convergence_tests_dir = os.path.join(runs_dir, 'convergence_tests')
        kpoint_tests_dir = os.path.join(convergence_tests_dir, 'kpoint_tests')
        encut_tests_dir = os.path.join(convergence_tests_dir, 'encut_tests')

        structure = Structure.from_file(poscar_path)

        if test_type == 'kpoints':
            test_kpoints()
        elif test_type == 'encut':
            test_encut()
        else:
            test_kpoints()
            test_encut()

        return {
            'status': 'success',
            'message': f'Generated VASP workflow of convergence tests of k-points in {kpoint_tests_dir} and energy cutoff in {encut_tests_dir}.',
            'kpoint_tests_dir': kpoint_tests_dir,
            'encut_tests_dir': encut_tests_dir,
        }
    
    except Exception as e:
        return {
            'status': 'error',
            'message': f'VASP convergence tests workflow generation failed: {str(e)}'
        }