from __future__ import annotations

import sys
from argparse import ArgumentParser
from ast import unparse
from pathlib import PurePosixPath

from . import __version__, fickle, tracing
from .analysis import Severity, check_safety
from .constants import EXIT_CLEAN, EXIT_ERROR, EXIT_UNSAFE
from .exception import ResourceExhaustionError
from .loader import scan_file, scan_zip_archive

DEFAULT_JSON_OUTPUT_FILE = "safety_results.json"

HF_RAW_PICKLE_EXTENSIONS = frozenset({".bin", ".pkl", ".pickle"})
HF_ZIP_PICKLE_EXTENSIONS = frozenset({".pt", ".pth"})
HF_PICKLE_EXTENSIONS = HF_RAW_PICKLE_EXTENSIONS | HF_ZIP_PICKLE_EXTENSIONS
HF_KNOWN_SAFE_EXTENSIONS = frozenset({".json", ".safetensors", ".onnx"})
HF_KNOWN_SAFE_FILES = frozenset({".gitattributes", ".gitignore", "README.md"})


def _scan_huggingface(
    repo_id: str,
    revision: str | None = None,
    token: str | None = None,
    json_output_path: str | None = None,
    print_results: bool = False,
) -> int:
    """Scan a HuggingFace Hub repository for potentially malicious pickle files.

    Args:
        repo_id: HuggingFace repository ID (e.g., 'bert-base-uncased')
        revision: Specific revision (branch, tag, or commit) to scan
        token: HuggingFace API token for private repositories
        json_output_path: Path to write JSON results
        print_results: Whether to print results to console

    Returns:
        EXIT_CLEAN (0), EXIT_UNSAFE (1), or EXIT_ERROR (2)
    """
    try:
        from huggingface_hub import HfApi, hf_hub_download

        from .polyglot import find_file_properties
    except ImportError:
        sys.stderr.write(
            "Error: huggingface_hub and torch are required for --huggingface scanning.\n"
            "Install with: pip install fickling[huggingface]\n"
        )
        return EXIT_ERROR

    api = HfApi(token=token)

    try:
        repo_info = api.repo_info(repo_id=repo_id, revision=revision, token=token)
    except Exception as e:  # noqa: BLE001 -- HF Hub may raise unpredictable HTTP errors
        sys.stderr.write(f"Error accessing HuggingFace repository '{repo_id}': {e!s}\n")
        return EXIT_ERROR

    if repo_info is None or not repo_info.siblings:
        if print_results:
            print(f"No files found in {repo_id}")
        return EXIT_CLEAN

    files_to_scan = [f.rfilename for f in repo_info.siblings]

    if not files_to_scan:
        if print_results:
            print(f"No scannable files found in {repo_id}")
        return EXIT_CLEAN

    if print_results:
        print(f"Scanning {len(files_to_scan)} file(s) in {repo_id}...")

    overall_safe = True
    failed = 0
    json_output = json_output_path or DEFAULT_JSON_OUTPUT_FILE

    for filename in files_to_scan:
        try:
            local_path = hf_hub_download(
                repo_id=repo_id,
                filename=filename,
                revision=revision,
                token=token,
            )
        except Exception as e:  # noqa: BLE001 -- HF Hub may raise unpredictable errors
            sys.stderr.write(f"Error downloading {filename}: {e!s}\n")
            overall_safe = False
            failed += 1
            continue

        ext = PurePosixPath(filename).suffix.lower()
        props = find_file_properties(local_path)
        is_zip = (
            props["is_torch_zip"] or props["is_standard_zip"] or ext in HF_ZIP_PICKLE_EXTENSIONS
        )
        if is_zip:
            if print_results:
                print(f"\n  Scanning: {filename}")
            member_results = scan_zip_archive(
                local_path, graceful=True, json_output_path=json_output
            )
            file_results = list(member_results.values())
        elif props["is_valid_pickle"] or ext in HF_RAW_PICKLE_EXTENSIONS:
            if print_results:
                print(f"\n  Scanning: {filename}")
            file_results = [scan_file(local_path, graceful=True, json_output_path=json_output)]
        elif props.get("is_7z") or props["is_tar"] or props["is_numpy_pickle"]:
            sys.stderr.write(
                f"Warning: '{filename}' detected as 7z/TAR/NumPy archive but scanning "
                f"these formats is not yet supported; skipping\n"
            )
            continue
        else:
            name = PurePosixPath(filename).name
            if ext not in HF_KNOWN_SAFE_EXTENSIONS and name not in HF_KNOWN_SAFE_FILES:
                sys.stderr.write(f"Warning: skipping '{filename}' (unknown extension '{ext}')\n")
            elif print_results:
                print(f"  Skipping safe file: {filename}")
            continue

        for result in file_results:
            if print_results:
                for ar in result.results:
                    result_str = ar.to_string()
                    if result_str:
                        print(f"    {result_str}")

            if not result.is_safe:
                overall_safe = False
                if print_results:
                    sys.stderr.write(f"    WARNING: {filename} may contain unsafe content!\n")

            if result.errors:
                for err in result.errors:
                    sys.stderr.write(f"    {err}\n")
                failed += 1

    if print_results:
        if failed > 0:
            sys.stderr.write(f"\nWARNING: {failed}/{len(files_to_scan)} file(s) failed to scan\n")
        if overall_safe:
            print(f"\n{repo_id}: No obvious safety issues detected")
        else:
            print(f"\n{repo_id}: Potentially unsafe content detected!")

    return EXIT_CLEAN if overall_safe else EXIT_UNSAFE


def main(argv: list[str] | None = None) -> int:
    if argv is None:
        argv = sys.argv

    parser = ArgumentParser(
        description="fickling is a static analyzer and interpreter for Python pickle data"
    )
    parser.add_argument(
        "PICKLE_FILE",
        type=str,
        nargs="?",
        default="-",
        help="path to the pickle file to either "
        "analyze or create (default is '-' for "
        "STDIN/STDOUT)",
    )
    options = parser.add_mutually_exclusive_group()
    options.add_argument(
        "--inject",
        "-i",
        type=str,
        default=None,
        help="inject the specified Python code to be run at the end of unpickling, "
        "and output the resulting pickle data",
    )
    parser.add_argument(
        "--inject-target",
        type=int,
        default=0,
        help="some machine learning frameworks stack multiple pickles into the same model file; "
        "this option specifies the index of the pickle file in which to inject the code from the "
        "`--inject` command (default is 0)",
    )
    options.add_argument("--create", "-c", type=str, default=None)
    parser.add_argument(
        "--run-last",
        "-l",
        action="store_true",
        help="used with --inject to have the injected code "
        "run after the existing pickling code in "
        "PICKLE_FILE (default is for the injected code "
        "to be run before the existing code)",
    )
    parser.add_argument(
        "--replace-result",
        "-r",
        action="store_true",
        help=(
            "used with --inject to replace the unpickling result of the code in PICKLE_FILE "
            "with the return value of the injected code. Either way, the preexisting pickling "
            "code is still executed."
        ),
    )
    options.add_argument(
        "--check-safety",
        "-s",
        action="store_true",
        help=(
            "test if the given pickle file is known to be unsafe. If so, exit with non-zero "
            "status. This test is not guaranteed correct; the pickle file may still be unsafe "
            "even if this check exits with code zero."
        ),
    )

    parser.add_argument(
        "--json-output",
        type=str,
        default=None,
        help="path to the output JSON file to store the analysis results from check-safety."
        f"If not provided, a default file named {DEFAULT_JSON_OUTPUT_FILE} will be used.",
    )

    parser.add_argument(
        "--print-results",
        "-p",
        action="store_true",
        help="Print the analysis results to the console when checking safety.",
    )

    parser.add_argument(
        "--trace",
        "-t",
        action="store_true",
        help="print a runtime trace while interpreting the input pickle file",
    )
    parser.add_argument("--version", "-v", action="store_true", help="print the version and exit")
    options.add_argument(
        "--huggingface",
        "--hf",
        type=str,
        default=None,
        metavar="REPO_ID",
        help="scan a model from HuggingFace Hub by repository ID (e.g., 'bert-base-uncased'). "
        "Requires huggingface_hub: pip install fickling[huggingface]",
    )
    parser.add_argument(
        "--hf-revision",
        type=str,
        default=None,
        help="specific revision (branch, tag, or commit) to scan from HuggingFace Hub",
    )
    parser.add_argument(
        "--hf-token",
        type=str,
        default=None,
        help="HuggingFace API token for private repositories. "
        "Prefer setting the HF_TOKEN environment variable to avoid token exposure in "
        "process listings",
    )

    args = parser.parse_args(argv[1:])

    if args.version:
        if sys.stdout.isatty():
            print(f"fickling version {__version__}")
        else:
            print(__version__)
        return EXIT_CLEAN

    # HuggingFace scanning mode
    if args.huggingface is not None:
        return _scan_huggingface(
            repo_id=args.huggingface,
            revision=args.hf_revision,
            token=args.hf_token,
            json_output_path=args.json_output,
            print_results=args.print_results,
        )

    if args.create is None:
        if args.PICKLE_FILE == "-":
            if hasattr(sys.stdin, "buffer") and sys.stdin.buffer is not None:
                file = sys.stdin.buffer
            else:
                file = sys.stdin
        else:
            file = open(args.PICKLE_FILE, "rb")
        try:
            stacked_pickled = fickle.StackedPickle.load(file, fail_on_decode_error=False)
        except fickle.PickleDecodeError as e:
            sys.stderr.write(f"Fickling failed to parse this pickle file. Error: {e!s}\n")
            if args.check_safety:
                sys.stderr.write(
                    "Parsing errors might be indicative of a maliciously crafted pickle file. DO NOT TRUST this file without performing further analysis!\n"
                )
                sys.stderr.write(
                    "\n(If this is a valid pickle file, please report the error at https://github.com/trailofbits/fickling)\n"
                )
            return EXIT_ERROR
        finally:
            file.close()

        if args.inject is not None:
            if args.inject_target >= len(stacked_pickled):
                sys.stderr.write(
                    f"Error: --inject-target {args.inject_target} is too high; there are only "
                    f"{len(stacked_pickled)} stacked pickle files in the input\n"
                )
                return EXIT_ERROR
            if hasattr(sys.stdout, "buffer") and sys.stdout.buffer is not None:
                buffer = sys.stdout.buffer
            else:
                buffer = sys.stdout
            for pickled in stacked_pickled[: args.inject_target]:
                pickled.dump(buffer)
            pickled = stacked_pickled[args.inject_target]
            if not isinstance(pickled[-1], fickle.Stop):
                sys.stderr.write(
                    "Warning: The last opcode of the input file was expected to be STOP, but was "
                    f"in fact {pickled[-1].info.name}"
                )
            pickled.insert_python_eval(
                args.inject,
                run_first=not args.run_last,
                use_output_as_unpickle_result=args.replace_result,
            )
            pickled.dump(buffer)
            for pickled in stacked_pickled[args.inject_target + 1 :]:
                pickled.dump(buffer)
        elif args.check_safety:
            was_safe = True
            json_output_path = args.json_output or DEFAULT_JSON_OUTPUT_FILE
            for pickled in stacked_pickled:
                safety_results = check_safety(pickled, json_output_path=json_output_path)

                # Print results if requested
                if args.print_results:
                    print(safety_results.to_string())

                if safety_results.severity > Severity.LIKELY_SAFE:
                    was_safe = False
                    if args.print_results:
                        sys.stderr.write(
                            "Warning: Fickling detected that the pickle file may be unsafe.\n\n"
                            "Do not unpickle this file if it is from an untrusted source!\n\n"
                        )

            return EXIT_CLEAN if was_safe else EXIT_UNSAFE

        else:
            var_id = 0
            for i, pickled in enumerate(stacked_pickled):
                interpreter = fickle.Interpreter(
                    pickled, first_variable_id=var_id, result_variable=f"result{i}"
                )
                try:
                    if args.trace:
                        trace = tracing.Trace(interpreter)
                        print(unparse(trace.run()))
                    else:
                        print(unparse(interpreter.to_ast()))
                except ResourceExhaustionError as e:
                    sys.stderr.write(
                        f"Error: {e}\n"
                        "This pickle file may contain an expansion attack. "
                        "Use --check-safety to analyze it.\n"
                    )
                    return 1
                var_id = interpreter.next_variable_id
    else:
        pickled = fickle.Pickled(
            [
                fickle.Global.create("__builtin__", "eval"),
                fickle.Mark(),
                fickle.Unicode(args.create.encode("utf-8")),
                fickle.Tuple(),
                fickle.Reduce(),
                fickle.Stop(),
            ]
        )
        if args.PICKLE_FILE == "-":
            file = sys.stdout
            if hasattr(file, "buffer") and file.buffer is not None:
                file = file.buffer
        else:
            file = open(args.PICKLE_FILE, "wb")
        try:
            pickled.dump(file)
        finally:
            file.close()

    return EXIT_CLEAN
