import os
import sys
import time

import httpx
import psutil
import pytest
import typer

# Ensure the current directory is in sys.path
sys.path.insert(0, os.getcwd())

from atopile.logging import LoggerForTest
from test.runner.common import (
    ORCHESTRATOR_URL_ENV,
    ClaimRequest,
    ClaimResponse,
    EventRequest,
    EventType,
)

# Set up test database logging for worker subprocess
# test_run_id is generated by the orchestrator and shared via env var
TEST_RUN_ID = os.environ.get("ATO_TEST_RUN_ID")
if TEST_RUN_ID:
    LoggerForTest.setup_logging(test_run_id=TEST_RUN_ID)

ORCHESTRATOR_URL = os.environ.get(ORCHESTRATOR_URL_ENV)

FBRK_TEST_WORKER_MEMORY_LIMIT = (
    int(os.environ.get("FBRK_TEST_WORKER_MEMORY_LIMIT_MB", 2 * 1024)) * 1024 * 1024
)

FBRK_TEST_TEST_TIMEOUT = int(os.environ.get("FBRK_TEST_TEST_TIMEOUT", 0))


def run_pytest(nodeid: str):
    pytest.main(
        [
            nodeid,
            "-vv",
            "-p",
            "test.runner.plugin",
            # Ensure co-located tests in `src/` are imported by package name,
            # not by path.
            "-p",
            "atopile.pytest_import_by_name",
            "--no-header",
            "--no-summary",
            "--disable-warnings",
            *["--timeout", "100"],
        ]
    )


def run_pytest_with_timeout(nodeid: str):
    import ctypes
    import threading

    timeout = FBRK_TEST_TEST_TIMEOUT
    if timeout == 0:
        run_pytest(nodeid)
        return

    main_thread_id = threading.current_thread().ident
    assert main_thread_id is not None  # Always valid for a running thread
    stop_watchdog = threading.Event()

    class ExceptionTestTimeout(BaseException): ...

    def _get_thread_stack(thread_id: int) -> str:
        """Capture the current stack trace of a thread by its ID."""
        import sys
        import traceback

        frames = sys._current_frames()
        if thread_id not in frames:
            return "<thread not found>"
        return "".join(traceback.format_stack(frames[thread_id]))

    def _watchdog():
        """
        Thread-based watchdog that injects TestTimeout into the main thread.
        Unlike signal handlers, this cannot be overridden by test code.
        """
        failures = 0
        last_printed_stack: str | None = None
        # Wait for initial timeout
        if stop_watchdog.wait(timeout):
            return  # Test completed normally

        # Test timed out - repeatedly inject exception until it sticks
        # (in case exceptions are caught/suppressed)
        while not stop_watchdog.is_set():
            if failures > 10:
                print("Test timed out too many times - exiting")
                sys.exit(0x59)  # random exit code
            wait_time = max(5 / (2**failures), 0.05)
            failures += 1
            # Format current time with milliseconds for higher precision
            t = time.time()
            time_s = (
                time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(t))
                + f".{int((t % 1) * 1000):03d}"
            )
            print(
                f"{time_s}: Test {nodeid} timed out - injecting exception"
                f" (waiting {wait_time:.2f}s)"
            )
            # Capture stack before injecting exception to see what's blocking
            stack_before = _get_thread_stack(main_thread_id)
            # one last extra check
            # Inject TestTimeout exception into the main thread
            ctypes.pythonapi.PyThreadState_SetAsyncExc(
                ctypes.c_ulong(main_thread_id),
                ctypes.py_object(ExceptionTestTimeout),
            )
            # Retry if the exception doesn't take
            if stop_watchdog.wait(timeout=wait_time):
                break

            # Exception didn't work - print where the thread was stuck
            stack_after = _get_thread_stack(main_thread_id)
            # Only print full stacks if they differ from the last printed one
            if stack_before != last_printed_stack:
                print(f"Stack when timeout was injected:\n{stack_before}")
                last_printed_stack = stack_before
            if stack_after != last_printed_stack:
                print(f"Stack after waiting (still blocked):\n{stack_after}")
                last_printed_stack = stack_after

    watchdog_thread = threading.Thread(target=_watchdog, daemon=True)
    watchdog_thread.start()

    try:
        run_pytest(nodeid)
        stop_watchdog.set()
    except ExceptionTestTimeout:
        stop_watchdog.set()
        pytest.fail(f"Test timed out after {timeout:.2f}s", pytrace=False)
    finally:
        stop_watchdog.set()
        watchdog_thread.join(timeout=0.1)


def main():
    if not ORCHESTRATOR_URL:
        print(f"{ORCHESTRATOR_URL_ENV} not set", file=sys.stderr)
        sys.exit(1)

    client = httpx.Client(timeout=10.0)
    pid = os.getpid()
    print(f"Worker {pid} started against {ORCHESTRATOR_URL}")

    # Keep session separate? Pytest reuses sys.modules.

    mem = psutil.Process().memory_info().rss

    try:
        while True:
            mem_usage = psutil.Process().memory_info().rss - mem
            if mem_usage > FBRK_TEST_WORKER_MEMORY_LIMIT:
                mem_usage_mb = mem_usage / 1024 / 1024
                print(
                    f"Worker {pid} memory usage exceeded limit: {mem_usage_mb:.2f}MB",
                    file=sys.stderr,
                )
                return
            try:
                request = ClaimRequest(pid=pid)
                resp = client.post(
                    f"{ORCHESTRATOR_URL}/claim", content=request.model_dump_json()
                )
                if resp.status_code != 200:
                    print(
                        f"Worker {pid} failed to claim test: {resp.status_code}",
                        file=sys.stderr,
                    )
                    time.sleep(1)
                    continue

                response = ClaimResponse.model_validate_json(resp.content)
                nodeid = response.nodeid

                if not nodeid:
                    print(f"Worker {pid} received no work, sleeping.")
                    time.sleep(1)
                    continue

                # Run pytest for this nodeid
                # We inject our http adapter plugin
                # Set terminal width for better debug output formatting
                os.environ["COLUMNS"] = "120"
                # Update the test name in the test logger for proper tracking
                LoggerForTest.update_test_name(nodeid)
                run_pytest_with_timeout(nodeid)

                # Flush logs after each test to ensure they're written to DB
                LoggerForTest.flush_all()

                # always exit
                # break

            except httpx.RequestError as e:
                print(f"Worker {pid} connection error: {e}", file=sys.stderr)
                time.sleep(1)
                # If the orchestrator is down, we might want to exit eventually
                continue
            except Exception as e:
                print(f"Worker {pid} unexpected error: {e}", file=sys.stderr)
                time.sleep(1)

    finally:
        print(f"Worker {pid} shutting down")
        # Flush any pending test logs to the database before exiting
        LoggerForTest.close_all()
        try:
            client.post(
                f"{ORCHESTRATOR_URL}/event",
                content=EventRequest(
                    type=EventType.EXIT, pid=pid, timestamp=time.time()
                ).model_dump_json(),
            )
        except Exception as e:
            print(f"Failed to send exit event: {e}", file=sys.stderr)
        client.close()


if __name__ == "__main__":
    typer.run(main)
