import argparse
import cbor2
import enum
import importlib
import json
import pathlib
import sys
import typing
import ast

import pycardano
from pycardano import PlutusData

import uplc
import uplc.ast

from . import compiler, builder, prelude, __version__, __copyright__
from .util import CompilerError, data_from_json
from .prelude import ScriptContext


class Command(enum.Enum):
    compile_pluto = "compile_pluto"
    compile = "compile"
    eval = "eval"
    parse = "parse"
    eval_uplc = "eval_uplc"
    build = "build"
    lint = "lint"


class Purpose(enum.Enum):
    spending = "spending"
    minting = "minting"
    rewarding = "rewarding"
    certifying = "certifying"
    any = "any"
    lib = "lib"


def plutus_data_from_json(annotation: typing.Type, x: dict):
    try:
        if annotation == int:
            return int(x["int"])
        if annotation == bytes:
            return bytes.fromhex(x["bytes"])
        if annotation is None:
            return None
        if isinstance(annotation, typing._GenericAlias):
            # Annotation is a List or Dict
            if annotation._name == "List":
                annotation_ann = annotation.__dict__["__args__"][0]
                return [plutus_data_from_json(annotation_ann, k) for k in x["list"]]
            if annotation._name == "Dict":
                annotation_key, annotation_val = annotation.__dict__["__args__"]
                return {
                    plutus_data_from_json(
                        annotation_key, d["k"]
                    ): plutus_data_from_json(annotation_val, d["v"])
                    for d in x["map"]
                }
        if issubclass(annotation, pycardano.PlutusData):
            return annotation.from_dict(x)
    except (KeyError, ValueError):
        raise ValueError(
            f"Annotation {annotation} does not match provided plutus datum {json.dumps(x)}"
        )


def plutus_data_from_cbor(annotation: typing.Type, x: bytes):
    try:
        if annotation in (int, bytes):
            return cbor2.loads(x)
        if annotation is None:
            return None
        if isinstance(annotation, typing._GenericAlias):
            # Annotation is a List or Dict
            if annotation._name == "List":
                annotation_ann = annotation.__dict__["__args__"][0]
                return [
                    plutus_data_from_cbor(annotation_ann, cbor2.dumps(k))
                    for k in cbor2.loads(x)
                ]
            if annotation._name == "Dict":
                annotation_key, annotation_val = annotation.__dict__["__args__"]
                return {
                    plutus_data_from_cbor(
                        annotation_key, cbor2.dumps(k)
                    ): plutus_data_from_cbor(annotation_val, v)
                    for k, v in cbor2.loads(x).items()
                }
        if issubclass(annotation, pycardano.PlutusData):
            return annotation.from_cbor(x)
    except (KeyError, ValueError):
        raise ValueError(
            f"Annotation {annotation} does not match provided plutus datum {x.hex()}"
        )


def check_params(
    command: Command,
    purpose: Purpose,
    validator_args,
    validator_params,
    force_three_params=False,
):
    if purpose == Purpose.any:
        # The any purpose does not do any checks. Use only if you know what you are doing
        return
    ret_type = validator_args[-1]
    # expect the validator to return None
    assert (
        ret_type is None or ret_type == prelude.Anything
    ), f"Expected contract to return None, but returns {ret_type}"

    num_onchain_params = 3 if purpose == Purpose.spending or force_three_params else 2
    onchain_params = validator_args[-1 - num_onchain_params : -1]
    param_types = validator_args[: -1 - num_onchain_params]
    required_onchain_parameters = 3 if purpose == Purpose.spending else 2
    assert (
        len(onchain_params) == required_onchain_parameters
    ), f"""\
{purpose.value.capitalize()} validator must expect {required_onchain_parameters} parameters at evaluation (on-chain), but was specified to have {len(onchain_params)}.
Make sure the validator expects parameters {'datum, ' if purpose == Purpose.spending else ''}redeemer and script context."""

    if command in (Command.eval, Command.eval_uplc):
        assert len(validator_params) == len(param_types) + len(
            onchain_params
        ), f"{purpose.value.capitalize()} validator expects {len(param_types) + len(onchain_params)} parameters for evaluation, but only got {len(validator_params)}."
    else:
        assert len(validator_params) == len(
            param_types
        ), f"{purpose.value.capitalize()} validator expects {len(onchain_params)} parameters at evaluation (on-chain) and {len(param_types)} parameters at compilation time, but got {len(validator_params)} during compilation."
    assert (
        onchain_params[-1] == ScriptContext
    ), f"Last parameter of the validator is always ScriptContext, but is {onchain_params[-1].__name__} here."


def perform_command(args):
    command = Command(args.command)
    purpose = Purpose(args.purpose)
    input_file = args.input_file if args.input_file != "-" else sys.stdin
    force_three_params = args.force_three_params
    constant_folding = args.constant_folding
    # read and import the contract
    with open(input_file, "r") as f:
        source_code = f.read()
    tmp_input_file = pathlib.Path("build").joinpath("__tmp_opshin.py")
    tmp_input_file.parent.mkdir(exist_ok=True)
    with tmp_input_file.open("w") as fp:
        fp.write(source_code)
    sys.path.append(str(pathlib.Path(tmp_input_file).parent.absolute()))
    sc = importlib.import_module(pathlib.Path(tmp_input_file).stem)
    sys.path.pop()
    # load the passed parameters if not a lib
    if purpose == Purpose.lib:
        assert not args.args, "Can not pass arguments to a library"
        parsed_params = []
    else:
        try:
            annotations = list(sc.validator.__annotations__.values())
        except AttributeError:
            raise AssertionError(
                f"Contract has no function called 'validator'. Make sure the compiled contract contains one function called 'validator' or {command.value} using `opshin {command.value} lib {str(input_file)}`."
            )
        if "return" not in sc.validator.__annotations__:
            annotations.append(prelude.Anything)
        parsed_params = []
        for i, (c, a) in enumerate(zip(annotations, args.args)):
            if a[0] == "{":
                try:
                    param_json = json.loads(a)
                except Exception as e:
                    raise ValueError(
                        f'Invalid parameter for contract passed at position {i}, expected json value, got "{a}". Did you correctly encode the value as json and wrap it in quotes?'
                    ) from e
                try:
                    param = plutus_data_from_json(c, param_json)
                except Exception as e:
                    raise ValueError(
                        f"Invalid parameter for contract passed at position {i}, expected type {c.__name__}."
                    ) from e
            else:
                try:
                    param_bytes = bytes.fromhex(a)
                except Exception as e:
                    raise ValueError(
                        "Expected hexadecimal CBOR representation of plutus datum but could not transform hex string to bytes."
                    ) from e
                try:
                    param = plutus_data_from_cbor(c, param_bytes)
                except Exception as e:
                    raise ValueError(
                        f"Invalid parameter for contract passed at position {i}, expected type {c.__name__}."
                    ) from e
            parsed_params.append(param)
        check_params(
            command,
            purpose,
            annotations,
            parsed_params,
            force_three_params,
        )

    if command == Command.eval:
        assert purpose != Purpose.lib, "Can not evaluate a library"
        print("Starting execution")
        print("------------------")
        try:
            ret = sc.validator(*parsed_params)
        except Exception as e:
            print(f"Exception of type {type(e).__name__} caused")
            ret = e
        print("------------------")
        print(ret)

    source_ast = compiler.parse(source_code, filename=input_file)

    if command == Command.parse:
        print("Parsed successfully.")
        return

    try:
        code = compiler.compile(
            source_ast,
            filename=input_file,
            force_three_params=force_three_params,
            validator_function_name="validator" if purpose != Purpose.lib else None,
            constant_folding=constant_folding,
            # do not remove dead code when compiling a library - none of the code will be used
            remove_dead_code=purpose != Purpose.lib,
        )
    except CompilerError as c:
        # Generate nice error message from compiler error
        if not isinstance(c.node, ast.Module):
            source_seg = ast.get_source_segment(source_code, c.node)
            start_line = c.node.lineno - 1
            end_line = start_line + len(source_seg.splitlines())
            source_lines = "\n".join(source_code.splitlines()[start_line:end_line])
            pos_in_line = source_lines.find(source_seg)
        else:
            start_line = 0
            pos_in_line = 0
            source_lines = source_code.splitlines()[0]

        overwrite_syntaxerror = (
            len("SyntaxError: ") * "\b" if command != Command.lint else ""
        )
        err = SyntaxError(
            f"""\
{overwrite_syntaxerror}{c.orig_err.__class__.__name__}: {c.orig_err}
Note that opshin errors may be overly restrictive as they aim to prevent code with unintended consequences.
""",
            (
                args.input_file,
                start_line + 1,
                pos_in_line,
                source_lines,
            )
            # we remove chaining so that users to not see the internal trace back,
        )
        err.orig_err = c.orig_err
        raise err from None

    if command == Command.compile_pluto:
        print(code.dumps())
        return
    code = code.compile()

    # apply parameters from the command line to the contract (instantiates parameterized contract!)
    code = code.term
    # UPLC lambdas may only take one argument at a time, so we evaluate by repeatedly applying
    for d in map(
        data_from_json,
        map(json.loads, (PlutusData.to_json(p) for p in parsed_params)),
    ):
        code = uplc.ast.Apply(code, d)
    code = uplc.ast.Program((1, 0, 0), code)

    if command == Command.compile:
        print(code.dumps())
        return

    if command == Command.build:
        if args.output_directory == "":
            if args.input_file == "-":
                print(
                    "Please supply an output directory if no input file is specified."
                )
                exit(-1)
            target_dir = pathlib.Path("build") / pathlib.Path(input_file).stem
        else:
            target_dir = pathlib.Path(args.output_directory)
        target_dir.mkdir(exist_ok=True, parents=True)
        artifacts = builder.generate_artifacts(builder._build(code))
        with (target_dir / "script.cbor").open("w") as fp:
            fp.write(artifacts.cbor_hex)
        with (target_dir / "script.plutus").open("w") as fp:
            fp.write(artifacts.plutus_json)
        with (target_dir / "script.policy_id").open("w") as fp:
            fp.write(artifacts.policy_id)
        with (target_dir / "mainnet.addr").open("w") as fp:
            fp.write(artifacts.mainnet_addr)
        with (target_dir / "testnet.addr").open("w") as fp:
            fp.write(artifacts.testnet_addr)

        print(f"Wrote script artifacts to {target_dir}/")
        return
    if command == Command.eval_uplc:
        print("Starting execution")
        print("------------------")
        assert isinstance(code, uplc.ast.Program)
        try:
            ret = uplc.dumps(uplc.eval(code))
        except Exception as e:
            print("An exception was raised")
            ret = e
        print("------------------")
        print(ret)


def parse_args():
    a = argparse.ArgumentParser(
        description="An evaluator and compiler from python into UPLC. Translate imperative programs into functional quasi-assembly."
    )
    a.add_argument(
        "command",
        type=str,
        choices=Command.__members__.keys(),
        help="The command to execute on the input file.",
        default="eval",
        nargs="?",
    )
    a.add_argument(
        "purpose",
        type=str,
        choices=Purpose.__members__.keys(),
        help="The intended script purpose. Determines the number of on-chain parameters "
        "(spending = 3, minting, rewarding, certifying = 2, any = no checks). "
        "This allows the compiler to check whether the correct amount of parameters was passed during compilation.",
        default="any",
        nargs="?",
    )
    a.add_argument(
        "input_file", type=str, help="The input program to parse. Set to - for stdin."
    )
    a.add_argument(
        "-o",
        "--output-directory",
        default="",
        type=str,
        help="The output directory for artefacts of the build command. Defaults to the filename of the compiled contract. of the compiled contract.",
    )
    a.add_argument(
        "--force-three-params",
        "--ftp",
        action="store_true",
        help="Enforces that the contract is always called with three virtual parameters on-chain. Enable if the script should support spending and other purposes.",
    )
    a.add_argument(
        "--constant-folding",
        "--cf",
        action="store_true",
        help="Enables experimental constant folding, including propagation and code execution.",
    )
    a.add_argument(
        "args",
        nargs="*",
        default=[],
        help="Input parameters for the validator (parameterizes the contract for compile/build). Either json or CBOR notation.",
    )
    a.add_argument(
        "--output-format-json",
        action="store_true",
        help="Changes the output of the Linter to a json format.",
    )
    a.add_argument(
        "--version",
        action="version",
        version=f"opshin {__version__} {__copyright__}",
    )
    return a.parse_args()


def main():
    args = parse_args()
    if Command(args.command) != Command.lint:
        perform_command(args)
    else:
        try:
            perform_command(args)
        except Exception as e:
            error_class_name = e.__class__.__name__
            message = str(e)
            if isinstance(e, SyntaxError):
                start_line = e.lineno
                pos_in_line = e.offset
                if hasattr(e, "orig_err"):
                    error_class_name = e.orig_err.__class__.__name__
                    message = str(e.orig_err)
            else:
                start_line = 1
                pos_in_line = 1
            if args.output_format_json:
                print(
                    convert_linter_to_json(
                        line=start_line,
                        column=pos_in_line,
                        error_class=error_class_name,
                        message=message,
                    )
                )
            else:
                print(
                    f"{args.input_file}:{start_line+1}:{pos_in_line}: {error_class_name}: {message}"
                )


def convert_linter_to_json(
    line: int,
    column: int,
    error_class: str,
    message: str,
):
    # output in lists
    return json.dumps(
        [
            {
                "line": line,
                "column": column,
                "error_class": error_class,
                "message": message,
            }
        ]
    )


if __name__ == "__main__":
    main()
