"""A runner for Markdown-based tests for ty"""
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "rich",
#     "watchfiles>=1.1.0",
# ]
#
# [tool.uv]
# exclude-newer = "7 days"
# ///

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import threading
from pathlib import Path
from typing import Final, Literal, assert_never

from rich.console import Console
from watchfiles import Change, watch

CRATE_NAME: Final = "ty_python_semantic"
CRATE_ROOT: Final = Path(__file__).resolve().parent
TY_VENDORED: Final = CRATE_ROOT.parent / "ty_vendored"
DIRS_TO_WATCH: Final = (
    CRATE_ROOT,
    TY_VENDORED,
    CRATE_ROOT.parent / "ty_test/src",
    CRATE_ROOT.parent / "mdtest/src",
)
MDTEST_DIR: Final = CRATE_ROOT / "resources" / "mdtest"
LINT_DOCS_DIR: Final = CRATE_ROOT / "resources" / "lint_docs"
MDTEST_SUITES: Final = (
    ("mdtest", MDTEST_DIR),
    ("lint_doc", LINT_DOCS_DIR),
)
SNAPSHOTS_DIR: Final = MDTEST_DIR / "snapshots"
MDTEST_README: Final = CRATE_ROOT / "resources" / "README.md"


class MDTestRunner:
    mdtest_executable: Path | None
    console: Console
    filters: list[str]
    enable_external: bool
    upgrade_lockfiles: bool
    update_snapshots: bool

    def __init__(
        self,
        filters: list[str] | None,
        enable_external: bool,
        upgrade_lockfiles: bool,
        update_snapshots: bool,
    ) -> None:
        self.mdtest_executable = None
        self.console = Console()
        self.filters = [f.removesuffix(".md") for f in (filters or [])]
        self.enable_external = enable_external
        self.upgrade_lockfiles = upgrade_lockfiles
        self.update_snapshots = update_snapshots

    def _run_cargo_test(self, *, message_format: Literal["human", "json"]) -> str:
        return subprocess.check_output(
            [
                "cargo",
                "test",
                "--package",
                CRATE_NAME,
                "--no-run",
                "--color=always",
                "--test=mdtest",
                "--message-format",
                message_format,
            ],
            cwd=CRATE_ROOT,
            env=dict(
                os.environ,
                CLI_COLOR="1",
                CARGO_PROFILE_DEV_OPT_LEVEL="0" if self.filters else "1",
                CARGO_PROFILE_DEV_DEBUG="line-tables-only",
            ),
            stderr=subprocess.STDOUT,
            text=True,
        )

    def _recompile_tests(
        self, status_message: str, *, message_on_success: bool = True
    ) -> bool:
        with self.console.status(status_message):
            # Run it with 'human' format in case there are errors:
            try:
                self._run_cargo_test(message_format="human")
            except subprocess.CalledProcessError as e:
                print(e.output)
                return False

            # Run it again with 'json' format to find the mdtest executable:
            try:
                json_output = self._run_cargo_test(message_format="json")
            except subprocess.CalledProcessError as _:
                # `cargo test` can still fail if something changed in between the two runs.
                # Here we don't have a human-readable output, so just show a generic message:
                self.console.print("[red]Error[/red]: Failed to compile tests")
                return False

            if json_output:
                self._get_executable_path_from_json(json_output)

        if message_on_success:
            self.console.print("[dim]Tests compiled successfully[/dim]")
        return True

    def _get_executable_path_from_json(self, json_output: str) -> None:
        for json_line in json_output.splitlines()[::-1]:
            try:
                data = json.loads(json_line)
            except json.JSONDecodeError:
                continue
            if data.get("target", {}).get("name") == "mdtest" and data["executable"]:
                self.mdtest_executable = Path(data["executable"])
                break
        else:
            raise RuntimeError(
                "Could not find mdtest executable after successful compilation"
            )

    def _run_mdtest(
        self, arguments: list[str] | None = None, *, capture_output: bool = False
    ) -> subprocess.CompletedProcess:
        assert self.mdtest_executable is not None

        arguments = arguments or []
        return subprocess.run(
            [self.mdtest_executable, *arguments],
            cwd=CRATE_ROOT,
            env=dict(
                os.environ,
                CLICOLOR_FORCE="1",
                INSTA_FORCE_PASS="1",
                INSTA_OUTPUT="none",
                MDTEST_EXTERNAL="1" if self.enable_external else "0",
                MDTEST_UPGRADE_LOCKFILES="1" if self.upgrade_lockfiles else "0",
                MDTEST_UPDATE_SNAPSHOTS="1" if self.update_snapshots else "0",
            ),
            capture_output=capture_output,
            text=True,
            check=False,
        )

    @staticmethod
    def _md_file_for_snapshot(snapshot_path: Path) -> tuple[str, Path] | None:
        """Given a deleted .snap.new path, find the source .md file it belongs to.

        Snapshot filenames follow the pattern:
            <md_filename>_-_<Section_title>_-_…_(<hash>).snap.new
        where <md_filename> may be truncated to 20 characters with a trailing
        Unicode ellipsis (U+2026).
        """
        snapshot_name = snapshot_path.name.removesuffix(".new").removesuffix(".snap")

        md_filename = snapshot_name.split("_-_", 1)[0]

        is_truncated = md_filename.endswith("\u2026")
        if is_truncated:
            md_filename = md_filename[:-1]

        for test_function, mdtest_dir in MDTEST_SUITES:
            for md_file in mdtest_dir.rglob("*.md"):
                if is_truncated:
                    if md_file.name.startswith(md_filename):
                        return test_function, md_file.relative_to(mdtest_dir)
                elif md_file.name == md_filename:
                    return test_function, md_file.relative_to(mdtest_dir)
        return None

    @staticmethod
    def _mdtest_for_path(path: Path) -> tuple[str, Path] | None:
        for test_function, mdtest_dir in MDTEST_SUITES:
            if path.is_relative_to(mdtest_dir):
                return test_function, path.relative_to(mdtest_dir)
        return None

    def _run_mdtests_for_file(self, test_function: str, markdown_file: Path) -> None:
        test_name = f"{test_function}::{markdown_file}"

        output = self._run_mdtest(["--exact", test_name], capture_output=True)

        if output.returncode == 0:
            if "running 0 tests\n" in output.stdout:
                self.console.log(
                    f"[yellow]Warning[/yellow]: No tests were executed with filter '{test_name}'"
                )
            else:
                self.console.print(
                    f"Test for [bold green]{markdown_file}[/bold green] succeeded"
                )
        else:
            self.console.print()
            self.console.rule(
                f"Test for [bold red]{markdown_file}[/bold red] failed",
                style="gray",
            )
            self._print_trimmed_cargo_test_output(
                output.stdout + output.stderr, test_name
            )

    def _print_trimmed_cargo_test_output(self, output: str, test_name: str) -> None:
        # Skip 'cargo test' boilerplate at the beginning:
        lines = output.splitlines()
        start_index = 0
        for i, line in enumerate(lines):
            if f"{test_name} stdout" in line:
                start_index = i
                break

        for line in lines[start_index + 1 :]:
            if "MDTEST_TEST_FILTER" in line:
                continue
            if line.strip() == "-" * 50:
                # Skip 'cargo test' boilerplate at the end
                break

            print(line)

    def watch(self):
        def keyboard_input() -> None:
            for _ in sys.stdin:
                # This is silly, but there is no other way to inject events into
                # the main `watch` loop. We use changes to the `README.md` file
                # as a trigger to re-run all mdtests:
                MDTEST_README.touch()

        input_thread = threading.Thread(target=keyboard_input, daemon=True)
        input_thread.start()

        self._recompile_tests("Compiling tests...", message_on_success=False)
        self._run_mdtest(self.filters)
        self.console.print("[dim]Ready to watch for changes...[/dim]")

        for changes in watch(*DIRS_TO_WATCH):
            new_mdtests: set[tuple[str, Path]] = set()
            changed_mdtests: set[tuple[str, Path]] = set()
            rejected_snapshot_mdtests: set[tuple[str, Path]] = set()
            rust_code_has_changed = False
            vendored_typeshed_has_changed = False

            for change, path_str in changes:
                path = Path(path_str)

                # See above: `README.md` changes trigger a full re-run of all tests
                if path == MDTEST_README:
                    self._run_mdtest(self.filters)
                    continue

                # When a pending snapshot (.snap.new) is rejected in a separate
                # process (e.g. `cargo insta review`), the file is deleted.
                # A common reason for rejecting a snapshot is that it was stale
                # (produced by an earlier test run and now outdated). Re-run the
                # relevant tests so the snapshot is regenerated from the current
                # state of the code.
                if (
                    change == Change.deleted
                    and path.name.endswith(".snap.new")
                    and path.is_relative_to(SNAPSHOTS_DIR)
                ):
                    md_file = self._md_file_for_snapshot(path)
                    if md_file is not None:
                        rejected_snapshot_mdtests.add(md_file)
                    continue

                match path.suffix:
                    case ".rs":
                        rust_code_has_changed = True
                    case ".pyi" if path.is_relative_to(TY_VENDORED):
                        vendored_typeshed_has_changed = True
                    case ".md":
                        pass
                    case _:
                        continue

                mdtest = self._mdtest_for_path(path)
                if mdtest is None:
                    continue

                match change:
                    case Change.added:
                        # When saving a file, some editors (looking at you, Vim) might first
                        # save the file with a temporary name (e.g. `file.md~`) and then rename
                        # it to the final name. This creates a `deleted` and `added` change.
                        # We treat those files as `changed` here.
                        if (Change.deleted, path_str) in changes:
                            changed_mdtests.add(mdtest)
                        else:
                            new_mdtests.add(mdtest)
                    case Change.modified:
                        changed_mdtests.add(mdtest)
                    case Change.deleted:
                        # No need to do anything when a Markdown test is deleted
                        pass
                    case _ as unreachable:
                        assert_never(unreachable)

            if rust_code_has_changed:
                if self._recompile_tests("Rust code has changed, recompiling tests..."):
                    self._run_mdtest(self.filters)
            elif vendored_typeshed_has_changed and self._recompile_tests(
                "Vendored typeshed has changed, recompiling tests..."
            ):
                self._run_mdtest(self.filters)

            for test_function, path in (
                new_mdtests | changed_mdtests | rejected_snapshot_mdtests
            ):
                self._run_mdtests_for_file(test_function, path)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="A runner for Markdown-based tests for ty"
    )
    parser.add_argument(
        "filters",
        nargs="*",
        help="Partial paths, e.g., 'loops/for.md' or 'invalid-argument-type.md'",
    )
    parser.add_argument(
        "--enable-external",
        "-e",
        action="store_true",
        help="Enable tests with external dependencies",
    )
    parser.add_argument(
        "--no-lockfile-upgrades",
        action="store_true",
        help="By default, lockfiles will be upgraded when dependency requirements in the Markdown test change."
        + " Set this flag to never upgrade any lockfiles.",
    )
    parser.add_argument(
        "--no-snapshot-updates",
        action="store_true",
        help="By default, inline snapshots will be updated automatically when they are stale. Set this flag to disable this.",
    )

    args = parser.parse_args()

    try:
        runner = MDTestRunner(
            filters=args.filters,
            enable_external=args.enable_external,
            upgrade_lockfiles=not args.no_lockfile_upgrades,
            update_snapshots=not args.no_snapshot_updates,
        )
        runner.watch()
    except KeyboardInterrupt:
        print()


if __name__ == "__main__":
    main()
