# flake8: noqa: E501
import argparse
import ast
import json

from evalscope.constants import EvalBackend, EvalType, JudgeStrategy, ModelTask


class ParseStrArgsAction(argparse.Action):

    def __call__(self, parser, namespace, values, option_string=None):
        assert isinstance(values, str), 'args should be a string.'

        # try json load first
        try:
            arg_dict = json.loads(values)
            setattr(namespace, self.dest, arg_dict)
            return
        except (json.JSONDecodeError, ValueError):
            pass

        # If JSON load fails, fall back to parsing as key=value pairs
        arg_dict = {}
        for arg in values.strip().split(','):
            key, value = map(str.strip, arg.split('=', 1))  # Use maxsplit=1 to handle multiple '='
            try:
                # Safely evaluate the value using ast.literal_eval
                arg_dict[key] = ast.literal_eval(value)
            except (ValueError, SyntaxError):
                # If ast.literal_eval fails, check if it's a boolean value
                value_lower = value.lower()
                if value_lower == 'true':
                    arg_dict[key] = True
                elif value_lower == 'false':
                    arg_dict[key] = False
                else:
                    # If not a boolean, keep the original string
                    arg_dict[key] = value

        setattr(namespace, self.dest, arg_dict)


def add_argument(parser: argparse.ArgumentParser):
    # yapf: disable
    # Model-related arguments
    parser.add_argument('--model', type=str, required=False, help='The model id on modelscope, or local model dir.')
    parser.add_argument('--model-id', type=str, required=False, help='The model id for model name in report.')
    parser.add_argument('--model-args', type=str, action=ParseStrArgsAction, help='The model args, should be a string.')
    parser.add_argument('--model-task', type=str, default=ModelTask.TEXT_GENERATION, choices=[ModelTask.TEXT_GENERATION, ModelTask.IMAGE_GENERATION], help='The model task for model id.')  # noqa: E501

    # Template-related arguments
    parser.add_argument('--chat-template', type=str, required=False, help='The custom jinja template for chat generation.')  # noqa: E501

    # Dataset-related arguments
    parser.add_argument('--datasets', type=str, nargs='+', required=False, help='Dataset id list, align to the module name in evalscope.benchmarks')  # noqa: E501
    parser.add_argument('--dataset-args', type=json.loads, default='{}', help='The dataset args, should be a json string.')  # noqa: E501
    parser.add_argument('--dataset-dir', help='The datasets dir.')
    parser.add_argument('--dataset-hub', help='The datasets hub.')

    # Generation configuration arguments
    parser.add_argument('--generation-config', type=str, action=ParseStrArgsAction, help='The generation config, should be a string.')  # noqa: E501

    # Evaluation-related arguments
    parser.add_argument('--eval-type', type=str,
                        choices=[
                            EvalType.CHECKPOINT,
                            EvalType.OPENAI_API,
                            EvalType.OPENAI_RESPONSES_API,
                            EvalType.ANTHROPIC_API,
                            EvalType.LITELLM,
                            EvalType.MOCK_LLM,
                            EvalType.TEXT2IMAGE,
                            EvalType.TEXT2SPEECH,
                            EvalType.IMAGE_EDITING,
                            EvalType.CUSTOM,
                        ],
                        help='Evaluation backend type. Auto-inferred from --api-url / --model when omitted.')  # noqa: E501
    parser.add_argument('--eval-backend', type=str, help='The evaluation backend to use.',
                        choices=[EvalBackend.NATIVE, EvalBackend.OPEN_COMPASS, EvalBackend.VLM_EVAL_KIT, EvalBackend.RAG_EVAL])  # noqa: E501
    parser.add_argument('--eval-config', type=str, required=False, help='The eval task config file path for evaluation backend.')  # noqa: E501
    parser.add_argument('--eval-batch-size', type=int, default=1, help='The batch size for evaluation.')
    parser.add_argument('--limit', type=float, default=None, help='Max evaluation samples num for each subset.')
    parser.add_argument('--repeats', type=int, default=1, help='Number of times to repeat the dataset items for k-metrics.')  # noqa: E501

    # Cache and working directory arguments
    parser.add_argument('--use-cache', type=str, help='Path to a previous output directory (e.g. outputs/20260519_120000) to resume from. Reuses cached predictions and reviews matched by sample_id.')  # noqa: E501
    parser.add_argument('--rerun-review', action='store_true', default=False, help='When --use-cache is set, force re-running the review/scoring step (deletes existing reviews cache) while still reusing prediction cache.')  # noqa: E501
    parser.add_argument('--work-dir', type=str, help='The root cache dir.')
    parser.add_argument('--no-timestamp', action='store_true', default=False, help='Do not add timestamp to work_dir to avoid overwriting previous results.')  # noqa: E501
    parser.add_argument('--enable-progress-tracker', action='store_true', default=False, help='Enable progress tracker.')

    # Debug and runtime mode arguments
    parser.add_argument('--ignore-errors', action='store_true', default=False, help='Ignore errors during evaluation.')
    parser.add_argument('--debug', action='store_true', default=False, help='Debug mode, will print information for debugging.')  # noqa: E501
    parser.add_argument('--seed', type=int, default=42, help='Random seed for reproducibility.')
    parser.add_argument('--api-key', type=str, default='EMPTY', help='The API key for the remote API model.')
    parser.add_argument('--api-url', type=str, default=None, help='The API url for the remote API model.')
    parser.add_argument('--timeout', type=float, default=None, help='[Deprecated] Use --generation-config timeout=... instead. Will be removed in v2.0.0.')  # noqa: E501
    parser.add_argument('--stream', action='store_true', default=None, help='[Deprecated] Use --generation-config stream=True instead. Will be removed in v2.0.0.')  # noqa: E501

    # LLMJudge arguments
    parser.add_argument('--judge-strategy', type=str, default=JudgeStrategy.AUTO, help='The judge strategy.')
    parser.add_argument('--judge-model-args', type=json.loads, default='{}', help='The judge model args, should be a json string.')  # noqa: E501
    parser.add_argument('--judge-worker-num', type=int, default=1, help='The number of workers for the judge model.')
    parser.add_argument('--analysis-report', action='store_true', default=False, help='Generate analysis report for the evaluation results using judge model.')  # noqa: E501
    parser.add_argument('--collect-perf', action=argparse.BooleanOptionalAction, default=True, help='Collect per-request performance metrics (latency, TTFT, token usage) during evaluation. TTFT requires streaming (--generation-config stream=True). Use --no-collect-perf to disable.')  # noqa: E501

    # Agent-config (native AgentLoop or external-agent bridge, distinguished
    # by ``mode``).  Native example:
    #   --agent-config '{"mode":"native","strategy":"react","max_steps":5}'
    # External example:
    #   --agent-config '{"mode":"external","framework":"claude-code","kwargs":{"model_name":"claude-opus-4-6"},"environment":"docker"}'
    parser.add_argument('--agent-config', type=json.loads, default=None, help='Agent configuration as a JSON object. Mode discriminated: {"mode":"native", ...} for AgentLoop, {"mode":"external","framework":"claude-code", ...} for external CLI bridge.')  # noqa: E501

    # Sandbox-related arguments
    parser.add_argument('--sandbox', type=json.loads, default=None, help='Unified sandbox configuration as a JSON string, e.g. \'{"enabled": true, "engine": "docker"}\'. Maps to TaskConfig.sandbox / SandboxTaskConfig.')  # noqa: E501
    parser.add_argument('--use-sandbox', action='store_true', default=False, help='[Deprecated] Use --sandbox instead. Whether to use sandbox for model evaluation.')  # noqa: E501
    parser.add_argument('--sandbox-type', type=str, default='docker', help='[Deprecated] Use --sandbox instead. The sandbox type to use (e.g., docker, volcengine).')  # noqa: E501
    parser.add_argument('--sandbox-manager-config', type=json.loads, default='{}', help='[Deprecated] Use --sandbox instead. The sandbox manager config, should be a json string.')  # noqa: E501
    # yapf: enable


def parse_args():
    parser = argparse.ArgumentParser(description='Run evaluation on benchmarks for LLMs.')
    add_argument(parser)

    args = parser.parse_args()
    return args
