import sys
import io
import json
import argparse
from typing import Dict, List

metrics = {}
metrics["has_variable"] = []
metrics["variable_eq"] = []


def code_execute(code: str, key: str) -> Dict[str, str]:
    """Executes code using Python's `exec`. Captures output
    (e.g., stdout, stderr) and reports that back.

    Returns:
        A dictionary containing model output and execution details.

    """
    env = {}
    old_stdout = sys.stdout
    old_stderr = sys.stderr

    with io.StringIO() as code_out:
        with io.StringIO() as code_err:
            sys.stdout = code_out
            sys.stderr = code_err
            executed = True
            execution_out = "ERROR"
            execution_err = None
            try:
                exec(code, env)
                execution_out = code_out.getvalue()
                execution_err = code_err.getvalue()
            except Exception as e:
                execution_out = "ERROR"
                execution_err = str(e)
                executed = False

            metrics[f"{key}_exec"] = executed
            metrics[f"{key}_out"] = execution_out
            metrics[f"{key}_err"] = execution_err

    sys.stdout = old_stdout
    sys.stderr = old_stderr

    return env


def run_code_execution(
    solution: str,
    model_code: str,
    metrics_out: str,
    required_variables: List[str],
    comparable_variables: List[str],
) -> None:
    """Runs code executation and execution-based evaluation.

    Arguments:
        solution: The reference code solution (i.e., what to match)
        model_code: Code produced by a model being tested.
        metrics_out: The path to print the metrics.
        required_variables: The variables that any mode produced code should contain.
        comparable_variables: Required variables that can be tested for equivalence.

    """

    ## execute solution
    try:
        env1 = code_execute(solution, "solution")
    except Exception as e:
        env1 = {}
        metrics["solution_exec"] = False
        metrics["solution_out"] = "ERROR"
        metrics["solution_err"] = str(e)

    ## execute model solution
    try:
        env2 = code_execute(model_code, "model")
    except Exception as e:
        env2 = {}
        metrics["model_exec"] = False
        metrics["model_out"] = "ERROR"
        metrics["model_err"] = str(e)

    ### check variables
    for variable in required_variables:
        metrics["has_variable"].append((variable, True if variable in env2 else False))

    ### check variable equivalence
    for variable in comparable_variables:
        try:
            self_eq = env1[variable] == env1[variable]
            if not isinstance(self_eq, bool):
                raise
        except:
            continue

        variable_eq = True
        try:
            env2[variable] == env2[variable]
            variable_eq = env1[variable] == env2[variable]
            if not isinstance(variable_eq, bool):
                continue
        except:
            variable_eq = False
            metrics["variable_eq"].append((variable, variable_eq))

    ## write outputs to metrics file
    with open(metrics_out, "w") as metrics_out:
        metrics_out.write(json.dumps(metrics))


parser = argparse.ArgumentParser()
parser.add_argument(
    "--code_file",
    help="Path to file containing all needed code and auxiliary data for executation evaluation (in JSON format)",
    type=str,
    required=True,
)


def main(argv) -> None:
    """Runs the main code execution and evaluation given a code artifact
    file.

    Arguments:
        argv: The raw command line arguments containing pointer to code data
            needed for evaluation.

    """
    args = parser.parse_args(sys.argv[1:])
    with open(args.code_file) as code_data:
        code_details = json.loads(code_data.read())

    run_code_execution(**code_details)


if __name__ == "__main__":
    main(sys.argv[1:])
