import argparse
import os
import signal
import subprocess
import sys
from datetime import datetime


class ProcessManager:
    """Manager for subprocess execution with proper signal handling."""

    def __init__(self):
        self.process = None
        self.original_handlers = {}

    def __enter__(self):
        """Context manager entry - setup signal handlers"""
        # Save original signal handlers
        self.original_handlers[signal.SIGINT] = signal.getsignal(signal.SIGINT)
        self.original_handlers[signal.SIGTERM] = signal.getsignal(signal.SIGTERM)

        # Register new signal handlers
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - restore original signal handlers"""
        # Restore original signal handlers
        for sig, handler in self.original_handlers.items():
            signal.signal(sig, handler)

    def _signal_handler(self, sig, frame):
        """Handle termination signals."""
        signal_name = 'SIGINT' if sig == signal.SIGINT else 'SIGTERM'
        print(f'\nReceived {signal_name}, cleaning up subprocess...')
        self.cleanup()
        sys.exit(0)

    def start_process(self, cmd):
        self.process = subprocess.Popen(cmd)
        return self.process

    def cleanup(self):
        if self.process and self.process.poll() is None:
            print('Terminating subprocess...')
            self.process.terminate()
            try:
                self.process.wait(timeout=5)
                print('Subprocess terminated successfully')
            except subprocess.TimeoutExpired:
                print('Subprocess did not terminate normally, forcing kill...')
                self.process.kill()
                self.process.wait()
                print('Subprocess killed')


def read_config():
    """Get configuration content from config file in script directory.

    Returns:
        str: Configuration file content, returns None if reading fails
    """
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(script_dir, 'config.py')

    # Read config file content
    try:
        with open(config_path, encoding='utf-8') as f:
            config_content = f.read()
        return config_content
    except FileNotFoundError:
        print(f'Error: Config file not found at {config_path}')
        return None
    except Exception as e:
        print(f'Error reading config file: {e}')
        return None


def update_datasets(config, datasets):
    """Update datasets part in config according to datasets list.

    Args:
        config (str): Original configuration content
        datasets (list[str]): List of dataset names to include
    Returns:
        str: Updated configuration content
    """
    datasets = list(datasets)
    if 'all' in datasets:
        # datasets part of the config file specifies all datasets, no need to update
        return config

    selected_datasets = []
    expanded_datasets = []
    for dataset in datasets:
        if dataset == 'aime':
            expanded_datasets.extend(['aime2024', 'aime2025'])
        else:
            expanded_datasets.append(dataset)
    datasets = []
    for dataset in expanded_datasets:
        if dataset not in datasets:
            datasets.append(dataset)

    if 'code' in datasets:
        selected_datasets.append('[LCBCodeGeneration_dataset]')
        datasets = [dataset for dataset in datasets if dataset != 'code']
    dataset_var_aliases = {
        'qmsum': 'LongBench_qmsum_datasets',
        'longbench_qmsum': 'LongBench_qmsum_datasets',
        'samsum': 'LongBench_samsum_datasets',
        'longbench_samsum': 'LongBench_samsum_datasets',
        'leval_meetingsumm': 'LEval_meetingsumm_datasets',
        'meetingsumm': 'LEval_meetingsumm_datasets',
        'meeting_summ': 'LEval_meetingsumm_datasets',
        'meetingbank': 'meetingbank_10k_datasets',
        'meetingbank_10k': 'meetingbank_10k_datasets',
    }
    for d in datasets:
        if d == 'arena':
            d = 'arenahard'
        selected_datasets.append(dataset_var_aliases.get(d, f'{d}_datasets'))
    selected_datasets = ' + '.join(selected_datasets)
    selected_datasets = f'datasets = {selected_datasets}'

    # replace datasets part in config
    start_tag = '# <dataset_replace_tag>'
    end_tag = '# </dataset_replace_tag>'

    start_index = config.find(start_tag)
    end_index = config.find(end_tag)

    if start_index == -1 or end_index == -1:
        raise ValueError('replace tag not found in config file')

    end_index += len(end_tag)
    replacement = f'{start_tag}\n{selected_datasets}\n{end_tag}'
    result = config[:start_index] + replacement + config[end_index:]
    return result


def get_model_name_from_server(server: str, tag: str) -> str:
    from openai import OpenAI
    try:
        client = OpenAI(api_key='YOUR_API_KEY', base_url=f'{server}/v1')
        model_name = client.models.list().data[0].id
        return model_name
    except Exception as e:
        raise RuntimeError(f'Failed to get model name from {tag}_server {server}: {e}')


def save_config(work_dir: str, config: str):
    """Save configuration content to a file in the specified directory.

    Args:
        work_dir (str): Directory to save the configuration file
        config (str): Configuration content to save
    """
    if not work_dir:
        return
    os.makedirs(work_dir, exist_ok=True)
    output_file = os.path.join(work_dir, 'config.py')
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(config)
    print(f'Config written to {output_file}')


def apply_runtime_overrides(config: str) -> str:
    """Apply environment-driven scalar overrides before OpenCompass parses config."""
    import re

    env_to_name = {
        'OPENCOMPASS_TARGET_TEMPERATURE': 'TARGET_TEMPERATURE',
        'OPENCOMPASS_TARGET_MAX_OUT_LEN': 'TARGET_MAX_OUT_LEN',
        'OPENCOMPASS_TARGET_QPS': 'TARGET_QPS',
        'OPENCOMPASS_TARGET_BATCH_SIZE': 'TARGET_BATCH_SIZE',
        'OPENCOMPASS_INFER_NUM_WORKER': 'INFER_NUM_WORKER',
        'OPENCOMPASS_INFER_MIN_TASK_SIZE': 'INFER_MIN_TASK_SIZE',
        'OPENCOMPASS_INFER_MAX_WORKERS': 'INFER_MAX_WORKERS',
        'OPENCOMPASS_EVAL_MAX_WORKERS': 'EVAL_MAX_WORKERS',
    }
    for env_name, config_name in env_to_name.items():
        value = os.environ.get(env_name)
        if value is None or value == '':
            continue
        if config_name in {'TARGET_TEMPERATURE', 'TARGET_QPS'}:
            literal = str(float(value))
        else:
            literal = str(int(value))
        pattern = rf'^{config_name}\s*=\s*.+$'
        replacement = f'{config_name} = {literal}'
        config, n = re.subn(pattern, replacement, config, count=1, flags=re.MULTILINE)
        if n == 0:
            raise ValueError(f'Could not apply OpenCompass runtime override for {config_name}')

    subjective_value = os.environ.get('OPENCOMPASS_SUBJECTIVE', '').strip().lower()
    if subjective_value:
        literal = 'True' if subjective_value in {'1', 'true', 'yes', 'on'} else 'False'
        pattern = r'^OPENCOMPASS_SUBJECTIVE\s*=\s*.+$'
        config, n = re.subn(pattern, f'OPENCOMPASS_SUBJECTIVE = {literal}', config, count=1, flags=re.MULTILINE)
        if n == 0:
            raise ValueError('Could not apply OpenCompass runtime override for OPENCOMPASS_SUBJECTIVE')

    ifeval_path = os.environ.get('OPENCOMPASS_IFEVAL_PATH', '').strip()
    if ifeval_path:
        config = config.replace("path='data/ifeval/input_data.jsonl'", f"path={ifeval_path!r}")
        config = config.replace('path="data/ifeval/input_data.jsonl"', f"path={ifeval_path!r}")
        config += (
            '\n\n# Runtime IFEval path override from lmdeploy/eval/eval.py\n'
            f'_runtime_ifeval_path = {ifeval_path!r}\n'
            'for _runtime_dataset in datasets:\n'
            "    if str(_runtime_dataset.get('abbr', '')).lower() == 'ifeval':\n"
            "        _runtime_dataset['path'] = _runtime_ifeval_path\n"
        )
    arena_hard_path = os.environ.get('OPENCOMPASS_ARENA_HARD_PATH', '').strip()
    if arena_hard_path:
        config = config.replace("path='./data/subjective/arena_hard'", f"path={arena_hard_path!r}")
        config = config.replace('path="./data/subjective/arena_hard"', f"path={arena_hard_path!r}")
        config += (
            '\n\n# Runtime Arena-Hard path override from lmdeploy/eval/eval.py\n'
            f'_runtime_arena_hard_path = {arena_hard_path!r}\n'
            'for _runtime_dataset in datasets:\n'
            "    if str(_runtime_dataset.get('abbr', '')).lower() == 'arenahard':\n"
            "        _runtime_dataset['path'] = _runtime_arena_hard_path\n"
        )

    test_range = os.environ.get('OPENCOMPASS_TEST_RANGE', '').strip()
    max_samples = os.environ.get('OPENCOMPASS_MAX_SAMPLES', '').strip()
    if test_range and max_samples:
        raise ValueError('Set only one of OPENCOMPASS_TEST_RANGE or OPENCOMPASS_MAX_SAMPLES')
    if max_samples:
        sample_count = int(max_samples)
        if sample_count <= 0:
            raise ValueError('OPENCOMPASS_MAX_SAMPLES must be positive')
        test_range = f'[:{sample_count}]'
    if test_range:
        config += (
            '\n\n# Runtime dataset range override from lmdeploy/eval/eval.py\n'
            f'_runtime_test_range = {test_range!r}\n'
            'for _runtime_dataset in datasets:\n'
            "    _runtime_dataset.setdefault('reader_cfg', {})['test_range'] = _runtime_test_range\n"
        )
    return config


def perform_evaluation(config, api_server, judger_server, mode, work_dir, reuse):
    """Perform model evaluation by opencompass.

    Args:
        config (str): Configuration content
        api_server (str): API server address for inference
        judger_server (str): Judger server address for evaluation
        mode (str): Running mode selection, options: infer, eval, all, config
        work_dir (str): Output directory for evaluation results. If not specified,
            config will not be saved and execution will not be performed.
        reuse (str): Whether to reuse existing results
    """
    direct_served_model = os.environ.get('OPENCOMPASS_SERVED_MODEL', '').strip()
    if mode in ['infer', 'all', 'config'] and api_server:
        served_model_name = direct_served_model or get_model_name_from_server(api_server, 'api')
        config = config.replace("SERVED_MODEL_PATH = ''", f"SERVED_MODEL_PATH = '{served_model_name}'")
    elif mode in ['infer', 'all']:
        raise ValueError('api-server is required when mode is infer or all')
    require_judger = os.environ.get('OPENCOMPASS_REQUIRE_JUDGER', '1').strip().lower() not in {
        '0', 'false', 'no', 'off'
    }
    direct_judger_model = os.environ.get('OPENCOMPASS_JUDGER_MODEL', '').strip()
    direct_judger_base = (
        os.environ.get('OPENCOMPASS_JUDGER_BASE_URL')
        or os.environ.get('OPENAI_BASE_URL')
        or 'https://api.openai.com'
    ).strip()
    if direct_judger_base.endswith('/v1'):
        direct_judger_base = direct_judger_base[:-3].rstrip('/')

    if mode in ['eval', 'all', 'config'] and judger_server:
        judger_model_name = get_model_name_from_server(judger_server, 'judger')
        config = config.replace("JUDGER_MODEL_PATH = ''", f"JUDGER_MODEL_PATH = '{judger_model_name}'")
    elif mode in ['eval', 'all', 'config'] and direct_judger_model:
        config = config.replace("JUDGER_MODEL_PATH = ''", f"JUDGER_MODEL_PATH = '{direct_judger_model}'")
        config = config.replace(
            "JUDGER_ADDR = 'http://<JUDGER_SERVER>'",
            f"JUDGER_ADDR = '{direct_judger_base}'",
        )
    elif mode in ['eval', 'all'] and require_judger:
        raise ValueError('judger-server is required when mode is eval or all')
    judger_api_key = os.environ.get('OPENCOMPASS_JUDGER_API_KEY') or os.environ.get('OPENAI_API_KEY')
    if judger_api_key:
        config = config.replace("key='YOUR_API_KEY'", f"key={judger_api_key!r}")
    config = apply_runtime_overrides(config)

    # write updated config to work_dir
    if work_dir:
        save_config(work_dir, config)
        if mode == 'config':
            return
    else:
        print(config)
        return

    # execute opencompass command
    cmd = ['opencompass', f'{work_dir}/config.py', '-m', mode, '-w', work_dir]
    if reuse:
        # reuse previous outputs & results. If reuse is a string, it indicates a specific timestamp.
        if reuse == 'latest':
            cmd.extend(['-r', str(reuse)])
        else:
            try:
                datetime.strptime(reuse, '%Y%m%d_%H%M%S')
                cmd.extend(['-r', str(reuse)])
            except ValueError as e:
                print(e)
                raise ValueError(f'Invalid reuse timestamp format: {reuse}. Expected format: YYYYMMDD_HHMMSS') from e
    try:
        print(f'Executing command: {" ".join(cmd)}')
        # result = subprocess.run(cmd, text=True, check=True)
        # return result
        with ProcessManager() as manager:
            process = manager.start_process(cmd)
            result = process.wait()
            if result != 0:
                raise subprocess.CalledProcessError(result, cmd)
            return subprocess.CompletedProcess(cmd, result)
    except Exception as e:
        print(f'Executing commanded failed with {e}')
        return


def main():
    parser = argparse.ArgumentParser(description='Perform model evaluation')
    parser.add_argument('task_name', type=str, help='The name of an evaluation task')
    parser.add_argument('-a', '--api-server', type=str, default='', help='API server address for inference')
    parser.add_argument('-j', '--judger-server', type=str, default='', help='Judger server address for evaluation')
    dataset_choices = [
        'aime',
        'aime2024',
        'aime2025',
        'gpqa',
        'gpqa_llada',
        'humaneval',
        'mbpp',
        'ifeval',
        'qmsum',
        'longbench_qmsum',
        'samsum',
        'longbench_samsum',
        'leval_meetingsumm',
        'meetingsumm',
        'meeting_summ',
        'meetingbank',
        'meetingbank_10k',
        'arenahard',
        'arena',
        'code',
        'mmlu_pro',
        'hle',
        'all',
    ]
    parser.add_argument('-d',
                        '--datasets',
                        nargs='+',
                        choices=dataset_choices,
                        default=['all'],
                        help=f"List of datasets. Available options: {', '.join(dataset_choices)}. "
                        'Use "all" to include all datasets.')
    parser.add_argument('-w',
                        '--work-dir',
                        type=str,
                        default='',
                        help='Output directory of evaluation. If not specified, outputs will not be saved.')
    parser.add_argument('-r',
                        '--reuse',
                        nargs='?',
                        type=str,
                        const='latest',
                        help='Reuse previous outputs & results, and run any missing jobs presented in the config. '
                        'If its argument is not specified, the latest results in the work_dir will be reused. '
                        'The argument should also be a specific timestamp, e.g. 20230516_144254')
    parser.add_argument('-m',
                        '--mode',
                        type=str,
                        help='Running mode selection. '
                        'all: complete pipeline including both inference and evaluation (default). '
                        'infer: only perform model inference to generate results. '
                        'eval: only evaluate previously generated results. '
                        'config: generate configuration files without execution.',
                        choices=['all', 'infer', 'eval', 'config'],
                        default='all')
    args = parser.parse_args()
    task_name = args.task_name
    api_server = args.api_server
    judger_server = args.judger_server
    datasets = args.datasets
    mode = args.mode
    work_dir = args.work_dir

    # Process server addresses
    if api_server and not api_server.startswith('http'):
        api_server = f'http://{api_server}'
    if judger_server and not judger_server.startswith('http'):
        judger_server = f'http://{judger_server}'

    # read config file
    config = read_config()

    # update task name in config
    config = config.replace("TASK_TAG = ''", f"TASK_TAG = '{task_name}'")

    # update datasets part of config according to args.datasets
    config = update_datasets(config, datasets)

    # update api_server part of config according to args.api_server
    if api_server:
        config = config.replace("API_SERVER_ADDR = 'http://<API_SERVER>'", f"API_SERVER_ADDR = '{api_server}'")
    if judger_server:
        # update judger_server part of config according to args.judger_server
        config = config.replace("JUDGER_ADDR = 'http://<JUDGER_SERVER>'", f"JUDGER_ADDR = '{judger_server}'")

    # perform evaluation
    perform_evaluation(config, api_server, judger_server, mode, work_dir, args.reuse)


if __name__ == '__main__':
    main()