# !/usr/bin/env python3

import os, sys, datetime, time
import tabulate
from pathlib import Path
from colorama import Fore, Style
from mp_api.client import MPRester
from openai import OpenAI
from importlib.metadata import version, PackageNotFoundError

def start_new_session():
    '''Set up a new session runs directory.'''
    base_dir = os.getcwd()
    main_dir = os.path.join(base_dir, 'masgent_projects')
    os.makedirs(main_dir, exist_ok=True)
    
    # Create a new runs directory with timestamp
    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
    runs_dir = os.path.join(main_dir, f'runs_{timestamp}')
    if not os.path.exists(runs_dir):
        os.makedirs(runs_dir, exist_ok=True)
    else:
        # Rare collision case, wait a second and try again
        time.sleep(1)
        timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        runs_dir = os.path.join(main_dir, f'runs_{timestamp}')
        os.makedirs(runs_dir, exist_ok=True)
    
    color_print(f'\n[Info] New Masgent session runs directory: {runs_dir}', 'green')
    os.environ['MASGENT_SESSION_RUNS_DIR'] = runs_dir

def global_commands():
    return [
        '',
        'AI    ->  Chat with the Masgent AI',
        'New   ->  Start a new session',
        'Back  ->  Return to previous menu',
        'Main  ->  Return to main menu',
        'Help  ->  Show available functions',
        'Exit  ->  Quit the Masgent',
    ]

def write_comments(file, file_type, comments):
    with open(file, 'r') as f:
        lines = f.readlines()

    if file_type.lower() in {'poscar', 'kpoints'}:
        lines[0] = comments + '\n'
    
    elif file_type.lower() in {'incar'}:
        lines.insert(0, f'{comments}\n')
    
    with open(file, 'w') as f:
        f.writelines(lines)

def get_color_map():
    return {
        'red': Fore.RED,
        'green': Fore.GREEN,
        'yellow': Fore.YELLOW,
        'blue': Fore.BLUE,
        'magenta': Fore.MAGENTA,
        'cyan': Fore.CYAN,
        'white': Fore.WHITE,
    }

def color_print(text, color='cyan'):
    '''Print text in specified color.'''
    color_map = get_color_map()
    chosen_color = color_map.get(color.lower(), Fore.CYAN)
    print(chosen_color + text + Style.RESET_ALL)

def color_input(text, color='cyan'):
    '''Input prompt in specified color.'''
    color_map = get_color_map()
    chosen_color = color_map.get(color.lower(), Fore.CYAN)
    return input(chosen_color + text + Style.RESET_ALL)

def load_system_prompts():
    # src/masgent/ai_mode/system_prompt.txt
    prompts_path = Path(__file__).resolve().parent / 'ai_mode' / 'system_prompt.txt'
    try:
        return prompts_path.read_text(encoding='utf-8')
    except Exception as e:
        return f'Error loading system prompts: {str(e)}'

def validate_openai_api_key(key):
    try:
        client = OpenAI(api_key=key)
        client.models.list()
        # color_print('[Info] OpenAI API key validated successfully.\n', 'green')
    except Exception as e:
        color_print('[Error] Invalid OpenAI API key. Exiting...\n', 'green')
        sys.exit(1)

def ask_for_openai_api_key():
    key = color_input('Enter your OpenAI API key: ', 'yellow').strip()
    if not key:
        color_print('[Error] OpenAI API key cannot be empty. Exiting...\n', 'green')
        sys.exit(1)
    
    validate_openai_api_key(key)

    os.environ['OPENAI_API_KEY'] = key

    save = color_input('Save this key to .env file for future? (y/n): ', 'yellow').strip().lower()
    base_dir = os.getcwd()
    env_path = os.path.join(base_dir, '.env')
    if save == 'y':
        with open(env_path, 'w') as f:
            f.write(f'OPENAI_API_KEY={key}\n')
        color_print(f'[Info] OpenAI API key saved to {env_path} file.\n', 'green')
    
def validate_mp_api_key(key):
    try:
        with MPRester(key, mute_progress_bars=True) as mpr:
            _ = mpr.materials.search(
                formula='Si',
                fields=['material_id']
            )
        # color_print('[Info] Materials Project API key validated successfully.\n', 'green')
    except Exception as e:
        color_print('[Error] Invalid Materials Project API key. Exiting...\n', 'green')
        sys.exit(1)
    
def ask_for_mp_api_key():
    key = color_input('Enter your Materials Project API key: ', 'yellow').strip()
    if not key:
        color_print('[Error] Materials Project API key cannot be empty. Exiting...\n', 'green')
        sys.exit(1)

    validate_mp_api_key(key)

    os.environ['MP_API_KEY'] = key

    save = color_input('Save this key to .env file for future? (y/n): ', 'yellow').strip().lower()
    base_dir = os.getcwd()
    env_path = os.path.join(base_dir, '.env')
    if save == 'y':
        with open(env_path, 'a') as f:
            f.write(f'MP_API_KEY={key}\n')
        color_print(f'[Info] Materials Project API key saved to {env_path} file.\n', 'green')

def print_banner():
    try:
        pkg_version = version('masgent')
    except PackageNotFoundError:
        pkg_version = 'dev'
    
    ascii_banner = rf'''
╔═════════════════════════════════════════════════════════════════════════╗
║                                                                         ║
║  ███╗   ███╗  █████╗  ███████╗  ██████╗  ███████╗ ███╗   ██╗ ████████╗  ║
║  ████╗ ████║ ██╔══██╗ ██╔════╝ ██╔════╝  ██╔════╝ ████╗  ██║ ╚══██╔══╝  ║
║  ██╔████╔██║ ███████║ ███████╗ ██║  ███╗ █████╗   ██╔██╗ ██║    ██║     ║
║  ██║╚██╔╝██║ ██╔══██║ ╚════██║ ██║   ██║ ██╔══╝   ██║╚██╗██║    ██║     ║
║  ██║ ╚═╝ ██║ ██║  ██║ ███████║ ╚██████╔╝ ███████╗ ██║ ╚████║    ██║     ║
║  ╚═╝     ╚═╝ ╚═╝  ╚═╝ ╚══════╝  ╚═════╝  ╚══════╝ ╚═╝  ╚═══╝    ╚═╝     ║
║                                                                         ║
║                                   MASGENT: Materials Simulation Agent   ║
║                                      Copyright (c) 2025 Guangchen Liu   ║
║                                                                         ║
║  Version:         {pkg_version:<52}  ║
║  Licensed:        MIT License                                           ║
║  Repository:      https://github.com/aguang5241/masgent                 ║
║  Citation:        Liu, G. et al. (2025), DOI:10.XXXX/XXXXX'             ║
║  Contact:         gliu4@wpi.edu                                         ║
║                                                                         ║
╚═════════════════════════════════════════════════════════════════════════╝
    '''
    color_print(ascii_banner, 'yellow')

def print_help():
    title = '\n Masgent - Available Commands and Functions '
    color_print(title, "green")

    headers = ['Index', 'Description']
    rows = [
        ['1', 'Density Functional Theory (DFT) Simulations'],
        
        ['1.1', 'Structure Preparation & Manipulation'],
        ['1.1.1', 'Generate POSCAR from chemical formula'],
        ['1.1.2', 'Convert POSCAR coordinates (Direct <-> Cartesian)'],
        ['1.1.3', 'Convert structure file formats (CIF, POSCAR, XYZ)'],
        ['1.1.4', 'Generate structures with defects (Vacancies, Interstitials, Substitutions)'],
        ['1.1.5', 'Generate supercells'],
        # ['1.1.6', 'Generate special quasirandom structures (SQS)'],
        # ['1.1.7', 'Generate surface slabs'],
        # ['1.1.8', 'Generate interface structures'],

        ['1.2', 'VASP Input File Preparation'],
        ['1.2.1', 'Prepare full VASP input files (INCAR, KPOINTS, POTCAR, POSCAR)'],
        ['1.2.2', 'Generate INCAR templates (relaxation, static, MD, etc.)'],
        ['1.2.3', 'Generate KPOINTS with specified accuracy'],
        ['1.2.4', 'Generate HPC job submission script'],
        # ['1.2.5', 'Generate standard VASP calculation workflows'],
        # ['1.2.5.1', 'Convergence testing (ENCUT, KPOINTS)'],
        # ['1.2.5.2', 'Equation of State (EOS)'],
        # ['1.2.5.3', 'Elastic constants'],

        ['1.3', 'VASP Output Analysis'],

        ['2', 'Machine Learning Potentials (MLP)'],

        ['3', 'Machine Learning Model Training & Evaluation'],
    ]
    
    table = tabulate.tabulate(rows, headers, tablefmt='fancy_grid')
    color_print(table, "green")
