import sys
import subprocess
# --- Dependency Check & Auto-Install ---
required_imports = [
    ("getpass", None),
    ("json", None),
    ("logging", None),
    ("os", None),
    ("threading", None),
    ("time", None),
    ("functools", None),
    ("multiprocessing", None),
    ("pathlib", None),
    ("typing", None),
    ("urllib.parse", None),
    ("pandas", "pandas"),
    ("requests", "requests"),
    ("tqdm", "tqdm"),
    ("pandas.io.formats.style", "pandas"),
]
for mod, pipname in required_imports:
    try:
        if "." in mod:
            __import__(mod.split(".")[0])
        else:
            __import__(mod)
    except ImportError:
        if pipname:
            print(f"Installing missing package: {pipname}")
            subprocess.check_call([sys.executable, "-m", "pip", "install", pipname])
        else:
            print(f"Module {mod} is a built-in or not installable via pip.")
# --- Script Description ---
"""
Autosimulator for WorldQuant BRAIN platform
- Timestamped logger
- Authentication with biometric check
- User-specified alpha JSON input
- Single/multi simulation mode
- Simulation worker: sends jobs, retries, saves locations
- Result worker: fetches results, saves to JSON
"""
import os
import sys
import time
import json
import threading
import logging
from datetime import datetime
from pathlib import Path
import requests
import getpass

# Platform specific imports
if sys.platform == 'win32':
    import msvcrt
else:
    import tty
    import termios
from ace_lib import (
    check_session_and_relogin,
    simulate_single_alpha,
    simulate_multi_alpha,
)

# --- Logger Setup ---
def setup_logger():
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_filename = f'autosim_{timestamp}.log'
    logger = logging.getLogger(f'autosim_{timestamp}')
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    fh = logging.FileHandler(log_filename)
    fh.setFormatter(formatter)
    logger.addHandler(fh)
    ch = logging.StreamHandler()
    ch.setFormatter(formatter)
    logger.addHandler(ch)
    return logger, log_filename

logger, log_filename = setup_logger()

# --- Authentication ---
def check_session_timeout(s):
    """
    Check if the current session has timed out.

    Args:
        s (SingleSession): The current session object.

    Returns:
        int: The number of seconds until the session expires, or 0 if the session has expired or an error occurred.
    """
    brain_api_url = os.environ.get("BRAIN_API_URL", "https://api.worldquantbrain.com")
    authentication_url = brain_api_url + "/authentication"
    try:
        result = s.get(authentication_url).json()["token"]["expiry"]
        logger.debug(f"Session (ID: {id(s)}) timeout check result: {result}")
        return result
    except Exception:
        logger.error(f"Session timeout check failed for session (ID: {id(s)})")
        return 0
def get_credentials():
    email = input("Email: ").strip()
    print("Password: ", end='', flush=True)
    password = []
    
    try:
        if sys.platform == 'win32':
            # Windows: Use msvcrt.getch()
            while True:
                char = msvcrt.getch()
                
                # Handle Enter key
                if char in [b'\r', b'\n']:
                    print()  # New line
                    break
                
                # Handle Backspace
                elif char == b'\x08':  # Backspace
                    if password:
                        password.pop()
                        # Move cursor back, print space, move cursor back again
                        print('\b \b', end='', flush=True)
                
                # Handle Ctrl+C
                elif char == b'\x03':  # Ctrl+C
                    print()
                    raise KeyboardInterrupt
                
                # Handle printable characters (ASCII)
                elif 32 <= ord(char) <= 126:  # Printable ASCII range
                    password.append(char.decode('ascii'))
                    print('*', end='', flush=True)
                
                # Handle extended characters
                else:
                    try:
                        decoded_char = char.decode('utf-8')
                        if decoded_char.isprintable():
                            password.append(decoded_char)
                            print('*', end='', flush=True)
                    except UnicodeDecodeError:
                        continue
        else:
            # Unix/macOS: Use tty and termios
            fd = sys.stdin.fileno()
            old_settings = termios.tcgetattr(fd)
            try:
                tty.setraw(fd)
                while True:
                    char = sys.stdin.read(1)
                    
                    # Handle Enter key
                    if char in ['\r', '\n']:
                        print('\r\n', end='', flush=True)
                        break
                    
                    # Handle Backspace
                    elif char in ['\x7f', '\x08']:
                        if password:
                            password.pop()
                            print('\b \b', end='', flush=True)
                    
                    # Handle Ctrl+C
                    elif char == '\x03':
                        print('\r\n', end='', flush=True)
                        raise KeyboardInterrupt
                    
                    # Handle printable characters
                    elif char.isprintable():
                        password.append(char)
                        print('*', end='', flush=True)
            finally:
                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
                
    except Exception as e:
        # Fallback to getpass
        print(f"\nError reading password: {e}")
        print("Falling back to getpass...")
        return (email, getpass.getpass())

    return (email, ''.join(password))

def authenticate():
    from ace_lib import SingleSession
    session = SingleSession()
    session.auth = get_credentials()
    brain_api_url = os.environ.get("BRAIN_API_URL", "https://api.worldquantbrain.com")
    r = session.post(brain_api_url + "/authentication")
    logger.debug(f"New session created (ID: {id(session)}) with authentication response: {r.status_code}, {r.json()}")
    if r.status_code == requests.status_codes.codes.unauthorized:
        if r.headers.get("WWW-Authenticate") == "persona":
            print(
                "Complete biometrics authentication and press any key to continue: \n"
                + r.url + "/persona?inquiry=" + r.headers.get("Location", "")
                + "\n"
            )
            input()
            session.post(r.headers.get("Location", r.url))
            while True:
                if session.post(r.headers.get("Location", r.url)).status_code != 201:
                    input(
                        "Biometrics authentication is not complete. Please try again and press any key when completed \n"
                    )
                else:
                    break
        else:
            logger.error("\nIncorrect email or password\n")
            return authenticate()
    return session

# --- User Input ---
MASTER_LOG_PATH = "autosim_master_log.json"

def update_master_log(input_json_path, latest_index):
    """
    Update the master log file with the latest successful index for the given input file name.
    """
    import os, json
    file_name = os.path.basename(input_json_path)
    log_data = {}
    # Read existing log if present
    if os.path.exists(MASTER_LOG_PATH):
        try:
            with open(MASTER_LOG_PATH, "r", encoding="utf-8") as f:
                log_data = json.load(f)
        except Exception:
            log_data = {}
    # Update with latest index
    log_data[file_name] = latest_index
    # Atomic write
    tmp_path = MASTER_LOG_PATH + ".tmp"
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(log_data, f, indent=2)
    os.replace(tmp_path, MASTER_LOG_PATH)
def get_user_json():
    import re
    while True:
        raw_path = input('Enter path to alpha JSON file: ').strip()
        json_path = re.sub(r'^["\']+|["\']+$', '', raw_path.strip())
        if os.path.exists(json_path):
            try:
                with open(json_path, 'r', encoding='utf-8') as f:
                    alpha_list = json.load(f)
                # Check master log for previous progress
                file_name = os.path.basename(json_path)
                start_index = 0
                if os.path.exists(MASTER_LOG_PATH):
                    try:
                        with open(MASTER_LOG_PATH, 'r', encoding='utf-8') as logf:
                            log_data = json.load(logf)
                        if file_name in log_data:
                            last_index = log_data[file_name]
                            print(f'Last time you simulated to position {last_index}.')
                            resp = input(f'Do you want to start from {last_index + 1}? (Y/n) Or enter another starting index: ').strip()
                            if resp.lower() in ['', 'y', 'yes']:
                                start_index = last_index + 1
                            elif resp.isdigit():
                                start_index = int(resp)
                            else:
                                print('Invalid input, starting from 0.')
                                start_index = 0
                    except Exception:
                        pass
                # Slice alpha_list to start from chosen index
                class AlphaList(list):
                    pass
                alpha_list = AlphaList(alpha_list[start_index:])
                alpha_list._start_index = start_index
                return alpha_list, json_path
            except Exception as e:
                logger.error(f'Error reading JSON file: {e}')
        else:
            logger.error(f'JSON file not found: {json_path}')
        print('Please enter a valid path to your alpha JSON file.')

def get_simulation_mode():
    mode = input('Select simulation mode (single/multi): ').strip().lower()
    if mode not in ['single', 'multi']:
        logger.error('Invalid mode. Choose "single" or "multi".')
        sys.exit(1)
    batch_size = None
    if mode == 'multi':
        while True:
            try:
                batch_size = int(input('Enter number of elements per multi-simulation batch (2-10): ').strip())
                if 2 <= batch_size <= 10:
                    break
                else:
                    print('Batch size must be between 2 and 10.')
            except Exception:
                print('Please enter a valid integer between 2 and 10.')
    return mode, batch_size

def get_retry_timeout():
    try:
        timeout = int(input('Enter retry timeout in seconds (default 60): ').strip())
        if timeout < 1:
            timeout = 60
    except Exception:
        timeout = 60
    return timeout

# --- Simulation Worker ---
def simulation_worker(session, alpha_list, mode, json_path, location_path, retry_timeout, batch_size=None):
    locations = {}
    # Initialize sent_count from user starting index (passed via alpha_list attribute if set)
    file_name = os.path.basename(json_path)
    sent_count = getattr(alpha_list, '_start_index', 0)
    while alpha_list:
        # Check session timeout before proceeding
        if check_session_timeout(session) == 0:
            logger.error('Session expired. Stopping simulation worker.')
            break
        session = check_session_and_relogin(session)
        # Prepare batch but do NOT pop yet
        if mode == 'single':
            batch = [alpha_list[0]]
        else:
            size = batch_size if batch_size else min(10, max(2, len(alpha_list)))
            batch = [alpha_list[i] for i in range(min(size, len(alpha_list)))]
        try:
            from ace_lib import start_simulation
            location = None
            while location is None:
                # Check session timeout before each send
                if check_session_timeout(session) == 0:
                    logger.error('Session expired. Stopping simulation worker.')
                    return
                if mode == 'single':
                    response = start_simulation(session, batch[0])
                    location = response.headers.get('Location')
                else:
                    response = start_simulation(session, batch)
                    location = response.headers.get('Location')
                if location is None:
                    logger.info(f'Simulation sent, location(s) saved: None')
                    logger.info(f'No location received, waiting {retry_timeout} seconds and retrying...')
                    time.sleep(retry_timeout)
            # Only pop/remove after location is valid
            if mode == 'single':
                alpha_list.pop(0)
                sent_count += 1
                update_master_log(json_path, sent_count - 1)
            else:
                for _ in range(len(batch)):
                    alpha_list.pop(0)
                sent_count += len(batch)
                update_master_log(json_path, sent_count - 1)
            locations[str(time.time())] = location
            with open(location_path, 'w', encoding='utf-8') as f:
                json.dump(locations, f, indent=2)
            # Do NOT overwrite the input JSON file
            logger.info(f'Simulation sent, location(s) saved: {location}')
        except Exception as e:
            logger.error(f'Simulation error: {e}. Retrying in {retry_timeout} seconds.')
            time.sleep(retry_timeout)

# --- Result Worker ---
def result_worker(session, location_path, result_path, poll_interval=30):
    results = {}
    from time import sleep
    while True:
        # Check session timeout before proceeding
        if check_session_timeout(session) == 0:
            logger.error('Session expired. Stopping result worker.')
            break
        session = check_session_and_relogin(session)
        if not os.path.exists(location_path):
            time.sleep(poll_interval)
            continue
        with open(location_path, 'r', encoding='utf-8') as f:
            locations = json.load(f)
        for loc_key, loc_val in locations.items():
            if loc_key in results:
                continue
            if not loc_val or not isinstance(loc_val, str) or not loc_val.startswith('http'):
                logger.error(f'Invalid or missing location for key {loc_key}: {loc_val}')
                continue
            try:
                # Check session timeout before each result fetch
                if check_session_timeout(session) == 0:
                    logger.error('Session expired. Stopping result worker.')
                    return
                simulation_progress_url = loc_val
                while True:
                    simulation_progress = session.get(simulation_progress_url)
                    retry_after = simulation_progress.headers.get("Retry-After", 0)
                    if float(retry_after) == 0:
                        break
                    logger.info(f"Sleeping for {retry_after} seconds for location {simulation_progress_url}")
                    sleep(float(retry_after))
                sim_json = simulation_progress.json()
                # Multi-simulation: check for children
                if "children" in sim_json and sim_json.get("status") == "COMPLETE":
                    child_results = {}
                    for child_id in sim_json["children"]:
                        child_url = f"https://api.worldquantbrain.com/simulations/{child_id}"
                        child_resp = session.get(child_url)
                        child_json = child_resp.json()
                        alpha_id = child_json.get("alpha")
                        if not alpha_id:
                            logger.error(f"No alpha_id found for child {child_id}")
                            child_results[child_id] = {"error": "No alpha_id found"}
                        else:
                            alpha = session.get(f"https://api.worldquantbrain.com/alphas/{alpha_id}")
                            child_results[child_id] = alpha.json()
                    results[loc_key] = {"multi_children": child_results}
                    logger.info(f"Multi-simulation results fetched for location {loc_val}")
                else:
                    # Single simulation
                    alpha_id = sim_json.get("alpha")
                    if not alpha_id:
                        logger.error(f"No alpha_id found for location {simulation_progress_url}")
                        results[loc_key] = {"error": "No alpha_id found"}
                    else:
                        alpha = session.get(f"https://api.worldquantbrain.com/alphas/{alpha_id}")
                        results[loc_key] = alpha.json()
                    logger.info(f"Result fetched for location {loc_val}")
                with open(result_path, 'w', encoding='utf-8') as f:
                    json.dump(results, f, indent=2)
            except Exception as e:
                logger.error(f'Error fetching result for {loc_val}: {e}')
        time.sleep(poll_interval)

# --- Main ---
def main():
    session = authenticate()
    alpha_list, json_path = get_user_json()
    mode, batch_size = get_simulation_mode()
    retry_timeout = get_retry_timeout()
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    location_path = f'autosim_locations_{timestamp}.json'
    result_path = f'autosim_results_{timestamp}.json'
    # Start workers
    sim_thread = threading.Thread(target=simulation_worker, args=(session, alpha_list, mode, json_path, location_path, retry_timeout, batch_size))
    res_thread = threading.Thread(target=result_worker, args=(session, location_path, result_path))
    sim_thread.start()
    res_thread.start()
    sim_thread.join()
    # Result worker runs until all locations processed
    logger.info('Simulation worker finished. Waiting for results...')
    res_thread.join()

if __name__ == '__main__':
    main()
