shell.py

import base64
import os
import re
import shlex
import sys

import src.var.flags as flags
from src.linter import LintRunner
from src.linter.context import LintOptions
from src.run.run import run
from src.run.source import read_source_file
from src.run.test_runner import run_tests
from src.var.constant import HELP_TEXT, VERSION
from src.var.keyword import FILE_FORMAT, TEST_FILE_EXTENSION


USE_DIRECTIVE_PATTERN = re.compile(
    r'^\s*@use\s+([A-Za-z_][A-Za-z0-9_]*)(?:\s+as\s+(.+?)|\s+(.+?))?\s*(?://.*)?$'
)
LEVEL_VALUES = {"error", "warning", "style", "security"}


def _extract_no_colors(args):
    filtered = []
    disabled = False
    for arg in args:
        if arg == "--nocolors":
            disabled = True
            continue
        filtered.append(arg)
    return filtered, disabled


def _strip_wrapping_quotes(value):
    if value is None:
        return None
    value = value.strip()
    if len(value) >= 2 and ((value[0] == '"' and value[-1] == '"') or (value[0] == "'" and value[-1] == "'")):
        return value[1:-1]
    return value


def collect_use_directives(source):
    directives = []
    for line_no, line in enumerate(source.splitlines(), start=1):
        match = USE_DIRECTIVE_PATTERN.match(line)
        if not match:
            continue
        name = match.group(1).lower()
        as_value = match.group(2)
        bare_value = match.group(3)
        has_as = as_value is not None
        raw_value = as_value if has_as else bare_value
        value = _strip_wrapping_quotes(raw_value) if raw_value is not None else None
        directives.append(
            {
                "name": name,
                "value": value,
                "has_as": has_as,
                "line": line_no,
            }
        )
    return directives


def apply_use_directives_to_lint_options(source, file_path, lint_options):
    merged = LintOptions(
        config_path=lint_options.config_path,
        level=lint_options.level,
        rules=lint_options.rules,
        fix=lint_options.fix,
        json_output=lint_options.json_output,
        failfast=lint_options.failfast,
    )

    for directive in collect_use_directives(source):
        name = directive["name"]
        value = directive["value"]
        has_as = directive["has_as"]

        if name == "save" and not str(file_path).lower().endswith(TEST_FILE_EXTENSION):
            raise ValueError("Directive '@use save' is available only in .test.omi files")

        if name == "json":
            merged.json_output = True
        elif name == "fix":
            merged.fix = True
        elif name == "failfast":
            merged.failfast = True
        elif name == "level":
            if not has_as or not value:
                raise ValueError("Directive '@use level' requires a value: @use level as <value>")
            level = value.lower()
            if level not in LEVEL_VALUES:
                raise ValueError("Directive '@use level' supports: error, warning, style, security")
            if merged.level is None:
                merged.level = level
        elif name == "rules":
            if not has_as or not value:
                raise ValueError("Directive '@use rules' requires a value: @use rules as <rule1,rule2>")
            parsed_rules = [item.strip() for item in value.split(",") if item.strip()]
            if merged.rules is None:
                merged.rules = parsed_rules
        elif name == "config":
            if merged.config_path is None and value:
                merged.config_path = value

    return merged


def uses_nolint(source):
    return any(directive["name"] == "nolint" for directive in collect_use_directives(source))


def apply_use_directives_to_test_flags(source, file_path, failfast, json_output, save_path):
    merged_failfast = failfast
    merged_json = json_output
    merged_save = save_path

    for directive in collect_use_directives(source):
        name = directive["name"]
        value = directive["value"]

        if name == "save" and not str(file_path).lower().endswith(TEST_FILE_EXTENSION):
            raise ValueError("Directive '@use save' is available only in .test.omi files")

        if name == "failfast":
            merged_failfast = True
        elif name == "json":
            merged_json = True
        elif name == "save":
            if merged_save is None and value:
                merged_save = value

    return merged_failfast, merged_json, merged_save


def parse_test_flags(args):
    failfast = False
    json_output = False
    save_path = None
    unknown = []

    i = 0
    while i < len(args):
        arg = args[i]

        if arg == "--failfast":
            failfast = True
        elif arg == "--json":
            json_output = True
        elif arg == "--save":
            if i + 1 < len(args) and not args[i + 1].startswith("--"):
                save_path = args[i + 1]
                i += 1
            else:
                save_path = None
        elif arg.startswith("--save="):
            value = arg.split("=", 1)[1].strip()
            save_path = value if value else None
        else:
            unknown.append(arg)

        i += 1

    return failfast, json_output, save_path, unknown


def parse_lint_flags(args):
    fix = False
    json_output = False
    failfast = False
    level = None
    rules = None
    config_path = None
    unknown = []

    i = 0
    while i < len(args):
        arg = args[i]

        if arg == "--fix":
            fix = True
        elif arg == "--json":
            json_output = True
        elif arg == "--failfast":
            failfast = True
        elif arg.startswith("--level="):
            value = arg.split("=", 1)[1].strip()
            level = value if value else None
        elif arg.startswith("--rules="):
            value = arg.split("=", 1)[1].strip()
            rules = [item.strip() for item in value.split(",") if item.strip()] if value else []
        elif arg == "--config":
            if i + 1 < len(args) and not args[i + 1].startswith("--"):
                config_path = args[i + 1]
                i += 1
            else:
                config_path = None
        elif arg.startswith("--config="):
            value = arg.split("=", 1)[1].strip()
            config_path = value if value else None
        else:
            unknown.append(arg)

        i += 1

    return (
        LintOptions(
            config_path=config_path,
            level=level,
            rules=rules,
            fix=fix,
            json_output=json_output,
            failfast=failfast,
        ),
        unknown,
    )


def parse_run_arguments(args):
    run_nolint = False
    lint_flag_tokens = []
    script_args = []
    parsing_flags = True

    i = 0
    while i < len(args):
        token = args[i]

        if parsing_flags and token == "--":
            script_args.extend(args[i + 1:])
            break
        if parsing_flags and token == "--lint":
            pass
        elif parsing_flags and token == "--nolint":
            run_nolint = True
        elif parsing_flags and token in ("--fix", "--json", "--failfast"):
            lint_flag_tokens.append(token)
        elif parsing_flags and (token.startswith("--level=") or token.startswith("--rules=") or token.startswith("--config=")):
            lint_flag_tokens.append(token)
        elif parsing_flags and token == "--config":
            lint_flag_tokens.append(token)
            if i + 1 < len(args) and not args[i + 1].startswith("--"):
                lint_flag_tokens.append(args[i + 1])
                i += 1
        elif parsing_flags and token.startswith("--"):
            lint_flag_tokens.append(token)
        else:
            parsing_flags = False
            script_args.append(token)

        i += 1

    lint_options, unknown_flags = parse_lint_flags(lint_flag_tokens)
    return run_nolint, lint_options, unknown_flags, script_args


def main(argv=None):
    args = argv if argv is not None else sys.argv[1:]
    args, no_colors_requested = _extract_no_colors(args)
    if no_colors_requested:
        flags.no_colors = True

    debug = ("--debug" in args) or ("-d" in args)

    if ("--version" in args) or ("-v" in args):
        print(f"Omi {VERSION}")
        return 0

    if ("--help" in args) or ("-h" in args):
        print(HELP_TEXT, end="")
        return 0

    cli_tokens = args

    if len(cli_tokens) >= 2 and cli_tokens[0] == "run":
        fn = cli_tokens[1]
        _, file_extension = os.path.splitext(fn)
        if file_extension not in FILE_FORMAT:
            print("Invalid file format (expected .omi)")
            return 1

        run_flags = cli_tokens[2:]
        run_flags, run_no_colors = _extract_no_colors(run_flags)
        if run_no_colors:
            flags.no_colors = True
        run_nolint, lint_options, unknown_flags, script_args = parse_run_arguments(run_flags)
        if unknown_flags:
            print(f"Unknown lint flag(s): {' '.join(unknown_flags)}")
            return 1

        try:
            script = read_source_file(fn)
        except Exception as e:
            print(f"Failed to load script \"{fn}\"\n{e}")
            return 1

        run_lint = not run_nolint and not uses_nolint(script)

        if run_lint:
            try:
                lint_options = apply_use_directives_to_lint_options(script, fn, lint_options)
            except ValueError as e:
                print(str(e))
                return 1

        result, error, file_flags = run(
            fn,
            script,
            lint_options=lint_options if run_lint else None,
            script_args=script_args,
            compact_lint_output=True,
        )
        if error:
            error_text = error.as_string()
            if error_text:
                print(error_text)
            return 1
        if (debug or file_flags.get("debug", False)) and result:
            if len(result.elements) == 1:
                print(repr(result.elements[0]))
            else:
                print(repr(result))
        return 0

    if len(cli_tokens) >= 2 and cli_tokens[0] == "test":
        target = cli_tokens[1]
        test_flag_tokens = cli_tokens[2:]
        test_flag_tokens, test_no_colors = _extract_no_colors(test_flag_tokens)
        if test_no_colors:
            flags.no_colors = True

        failfast, json_output, save_path, unknown_flags = parse_test_flags(test_flag_tokens)

        if unknown_flags:
            print(f"Unknown test flag(s): {' '.join(unknown_flags)}")
            return 1

        if os.path.isfile(target) and not target.lower().endswith(TEST_FILE_EXTENSION):
            print("RTError: Test files must have .test.omi extension")
            return 1
        try:
            if os.path.isfile(target):
                test_source = read_source_file(target)
                failfast, json_output, save_path = apply_use_directives_to_test_flags(
                    test_source,
                    target,
                    failfast,
                    json_output,
                    save_path,
                )

            exit_code = run_tests(
                target,
                failfast=failfast,
                json_output=json_output,
                save_path=save_path,
            )
        except ValueError as e:
            print(str(e))
            return 1
        except Exception as e:
            print(f"Failed to run tests for '{target}'\n{e}")
            return 1
        return exit_code

    if len(cli_tokens) >= 2 and cli_tokens[0] == "lint":
        target = cli_tokens[1]
        lint_flag_tokens, lint_no_colors = _extract_no_colors(cli_tokens[2:])
        if lint_no_colors:
            flags.no_colors = True
        lint_options, unknown_flags = parse_lint_flags(lint_flag_tokens)

        if unknown_flags:
            print(f"Unknown lint flag(s): {' '.join(unknown_flags)}")
            return 1

        if os.path.isfile(target) and not (
            target.lower().endswith(".omi") or target.lower().endswith(TEST_FILE_EXTENSION)
        ):
            print("Invalid file format (expected .omi or .test.omi)")
            return 1

        try:
            if os.path.isfile(target):
                lint_source = read_source_file(target)
                lint_options = apply_use_directives_to_lint_options(lint_source, target, lint_options)

            runner = LintRunner(
                config_path=lint_options.config_path,
                level=lint_options.level,
                rules=lint_options.rules,
                fix=lint_options.fix,
                json_output=lint_options.json_output,
                failfast=lint_options.failfast,
            )
            result = runner.lint_path(target)
            if lint_options.json_output:
                print(result.report.to_json())
            else:
                print(result.report.to_text())
            return result.exit_code
        except ValueError as e:
            print(str(e))
            return 1
        except Exception as e:
            print(f"Failed to run linter for '{target}'\n{e}")
            return 1

    while True:
        try:
            text = input("OmiShell >>> ")
            if text.strip() == "":
                continue

            _x = bytes.fromhex("676f6f6e").decode()
            if text.strip() == _x:
                try:
                    shell_file = os.path.join(os.path.dirname(__file__), "src", "nodes", "shell.py")
                    with open(shell_file, "r") as f:
                        encoded_content = f.read()
                    decoded_content = base64.b64decode(encoded_content).decode()
                    print(decoded_content)
                except Exception as e:
                    print(f"Error: {e}")
                continue

            if text.strip().startswith("run "):
                try:
                    command_tokens = shlex.split(text.strip())
                except ValueError as e:
                    print(f"Invalid run command: {e}")
                    continue

                if len(command_tokens) < 2:
                    print("Usage: run <file.omi> [flags]")
                    continue

                fn = command_tokens[1]
                run_flags, run_no_colors = _extract_no_colors(command_tokens[2:])
                if run_no_colors:
                    flags.no_colors = True

                run_nolint, lint_options, unknown_flags, script_args = parse_run_arguments(run_flags)
                if unknown_flags:
                    print(f"Unknown lint flag(s): {' '.join(unknown_flags)}")
                    continue

                _, file_extension = os.path.splitext(fn)

                if file_extension not in FILE_FORMAT:
                    print("Invalid file format (expected .omi)")
                    continue

                try:
                    script = read_source_file(fn)
                except Exception as e:
                    print(f"Failed to load script \"{fn}\"\n{e}")
                    continue

                run_lint = not run_nolint and not uses_nolint(script)

                if run_lint:
                    try:
                        lint_options = apply_use_directives_to_lint_options(script, fn, lint_options)
                    except ValueError as e:
                        print(str(e))
                        continue

                result, error, file_flags = run(
                    fn,
                    script,
                    lint_options=lint_options if run_lint else None,
                    script_args=script_args,
                    compact_lint_output=True,
                )
                if error:
                    error_text = error.as_string()
                    if error_text:
                        print(error_text)
                elif (debug or file_flags.get("debug", False)) and result:
                    if len(result.elements) == 1:
                        print(repr(result.elements[0]))
                    else:
                        print(repr(result))

                if flags.repl_output_emitted and not flags.repl_output_ended_with_newline:
                    print()

                continue

            if text.strip().startswith("test "):
                try:
                    command_tokens = shlex.split(text.strip())
                except ValueError as e:
                    print(f"Invalid test command: {e}")
                    continue

                if len(command_tokens) < 2:
                    print("Usage: test <file.test.omi|directory> [--failfast] [--json] [--save[=path]]")
                    continue

                target = command_tokens[1]
                failfast, json_output, save_path, unknown_flags = parse_test_flags(command_tokens[2:])
                if unknown_flags:
                    print(f"Unknown test flag(s): {' '.join(unknown_flags)}")
                    continue

                if os.path.isfile(target) and not target.lower().endswith(TEST_FILE_EXTENSION):
                    print("RTError: Test files must have .test.omi extension")
                    continue

                try:
                    run_tests(
                        target,
                        failfast=failfast,
                        json_output=json_output,
                        save_path=save_path,
                    )
                except ValueError as e:
                    print(str(e))
                except Exception as e:
                    print(f"Failed to run tests for '{target}'\n{e}")

                if flags.repl_output_emitted and not flags.repl_output_ended_with_newline:
                    print()

                continue

            result, error, _ = run("<stdin>", text)

            if error:
                error_text = error.as_string()
                if error_text:
                    print(error_text)
            elif debug and result:
                if len(result.elements) == 1:
                    print(repr(result.elements[0]))
                else:
                    print(repr(result))

            if flags.repl_output_emitted and not flags.repl_output_ended_with_newline:
                print()

        except KeyboardInterrupt:
            return 0


if __name__ == "__main__":
    raise SystemExit(main())