import asyncio
import json
import os
import re
import secrets
import subprocess
import time
import uuid
from contextlib import contextmanager, redirect_stdout, suppress
from datetime import datetime, timedelta, timezone
from io import StringIO
from typing import (
    AsyncIterator,
    Dict,
    Generator,
    Iterator,
    List,
    Optional,
    Tuple,
    Union,
)

import httpx
import pytest
from fastapi import Request, WebSocket
from httpx import HTTPStatusError
from isolate.backends.common import active_python
from openapi_fal_rest.api.applications import app_metadata
from openapi_fal_rest.client import Client
from pydantic import BaseModel
from pydantic import __version__ as pydantic_version
from websockets.sync import client as ws_client

import fal
import fal.api as api
from fal import apps
from fal.api.deploy import User, _get_user
from fal.app import REQUEST_ID_KEY, AppClient, AppClientError, wrap_app
from fal.auth import key_credentials
from fal.container import ContainerImage
from fal.exceptions import (
    AppException,
    FalServerlessException,
    FieldException,
    RequestCancelledException,
)
from fal.exceptions._cuda import _CUDA_OOM_MESSAGE, _CUDA_OOM_STATUS_CODE
from fal.ref import get_current_app
from fal.sdk import RunnerState, get_default_credentials
from fal.toolkit.utils.endpoint import cancel_on_disconnect
from fal.workflows import Workflow


@pytest.fixture(scope="module")
def rest_client() -> Generator[Client, None, None]:
    client = api.client.SyncServerlessClient()
    yield client._create_rest_client()


class Input(BaseModel):
    lhs: int
    rhs: int
    wait_time: int = 0


class StatefulInput(BaseModel):
    value: int


class FieldInput(BaseModel):
    value: Union[int, str, float]

    class Config:
        smart_union = True


class Output(BaseModel):
    result: int


actual_python = active_python()


def _auth_headers() -> Dict[str, str]:
    key_creds = key_credentials()
    if not key_creds:
        return {}
    key_id, key_secret = key_creds
    return {"Authorization": f"Key {key_id}:{key_secret}"}


def git_revision_short_hash() -> str:
    return (
        subprocess.check_output(["git", "rev-parse", "--short", "HEAD"])
        .decode("ascii")
        .strip()
    )


@fal.function(
    keep_alive=60,
    machine_type=["S", "M"],
    serve=True,
    max_concurrency=1,
    requirements=[f"pydantic=={pydantic_version}"],
)
def addition_app(input: Input) -> Output:
    print("starting...")
    for _ in range(input.wait_time):
        print("sleeping...")
        time.sleep(1)

    return Output(result=input.lhs + input.rhs)


nomad_addition_app = addition_app.on(_scheduler="nomad")


@fal.function(
    kind="container",
    image=ContainerImage.from_dockerfile_str(
        f"FROM python:{actual_python}-slim\n# {git_revision_short_hash()}",
    ),
    keep_alive=60,
    machine_type="S",
    serve=True,
    max_concurrency=1,
)
def container_addition_app(input: Input) -> Output:
    print("starting...")
    for _ in range(input.wait_time):
        print("sleeping...")
        time.sleep(1)

    return Output(result=input.lhs + input.rhs)


@fal.function(
    kind="container",
    image=ContainerImage.from_dockerfile_str(
        f"FROM python:{actual_python}-slim\n# {git_revision_short_hash()}",
    ),
    keep_alive=60,
    machine_type="S",
    serve=True,
    max_concurrency=1,
    force_env_build=False,
)
def container_cache_enabled_app(input: Input) -> Output:
    print("starting...")
    for _ in range(input.wait_time):
        print("sleeping...")
        time.sleep(1)

    return Output(result=input.lhs + input.rhs)


@fal.function(
    kind="container",
    image=ContainerImage.from_dockerfile_str(
        f"FROM python:{actual_python}-slim\n# {git_revision_short_hash()}",
    ),
    keep_alive=60,
    machine_type="S",
    serve=True,
    max_concurrency=1,
    force_env_build=True,
)
def container_no_cache_app(input: Input) -> Output:
    print("starting...")
    for _ in range(input.wait_time):
        print("sleeping...")
        time.sleep(1)

    return Output(result=input.lhs + input.rhs)


@fal.function(
    kind="container",
    image=ContainerImage.from_dockerfile_str(
        f"""FROM python:{actual_python}-slim\n# {git_revision_short_hash()}
ARG OUTPUT="built incorrectly"
ENV OUTPUT="$OUTPUT"
""",
        build_args={"OUTPUT": "built with build args"},
    ),
    keep_alive=60,
    machine_type="S",
    serve=True,
    max_concurrency=1,
)
def container_build_args_app() -> str:
    import os

    return os.environ["OUTPUT"]


@fal.function(
    keep_alive=300,
    requirements=["fastapi", "uvicorn", "pydantic==1.10.18"],
    machine_type="S",
    max_concurrency=1,
    max_multiplexing=30,
    exposed_port=8000,
)
def calculator_app():
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware
    from uvicorn import run

    app = FastAPI()

    def _wait(wait_time: int):
        print("starting...")
        for _ in range(wait_time):
            print("sleeping...")
            time.sleep(1)

    @app.post("/add")
    def add(input: Input) -> Output:
        _wait(input.wait_time)
        return Output(result=input.lhs + input.rhs)

    @app.post("/subtract")
    def subtract(input: Input) -> Output:
        _wait(input.wait_time)
        return Output(result=input.lhs - input.rhs)

    app.add_middleware(
        CORSMiddleware,
        allow_credentials=True,
        allow_headers=("*"),
        allow_methods=("*"),
        allow_origins=("*"),
    )

    run(app, host="0.0.0.0", port=8080)


class StatefulAdditionApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "S"

    async def setup(self):
        self.counter = 0

    @fal.endpoint("/reset")
    def reset(self) -> Output:
        self.counter = 0
        return Output(result=self.counter)

    @fal.endpoint("/increment")
    def increment(self, input: StatefulInput) -> Output:
        self.counter += input.value
        return Output(result=self.counter)

    @fal.endpoint("/decrement")
    def decrement(self, input: StatefulInput) -> Output:
        self.counter -= input.value
        return Output(result=self.counter)


class SleepInput(BaseModel):
    wait_time: int


class SleepOutput(BaseModel):
    slept: bool = True


class SleepApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "XS"

    @fal.endpoint("/")
    async def sleep(self, input: SleepInput) -> SleepOutput:
        for _ in range(input.wait_time):
            print("sleeping...", flush=True)
            await asyncio.sleep(1)
        return SleepOutput(slept=True)


class QueueBlockingApp(fal.App, keep_alive=300, max_concurrency=1, max_multiplexing=1):
    """
    App for testing start_timeout with queue blocking.

    With max_concurrency=1 and max_multiplexing=1, only ONE request can be
    processed at a time. Additional requests must wait in the queue.
    """

    machine_type = "XS"

    @fal.endpoint("/")
    async def sleep(self, input: SleepInput) -> SleepOutput:
        for i in range(input.wait_time):
            print(f"sleeping {i + 1}/{input.wait_time}...", flush=True)
            await asyncio.sleep(1)
        return SleepOutput(slept=True)


class ExceptionApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "XS"

    @fal.endpoint("/fail")
    def fail(self) -> Output:
        raise Exception("this app is designed to fail!")

    @fal.endpoint("/app-exception")
    def app_exception(self) -> Output:
        raise AppException(message="this app is designed to fail", status_code=401)

    @fal.endpoint("/field-exception")
    def field_exception(self, input: Input) -> Output:
        raise FieldException(
            field="rhs",
            message="rhs must be an integer",
        )

    @fal.endpoint("/field-exception-units")
    def field_exception_units(self, input: FieldInput) -> Output:
        raise FieldException(
            field="value",
            message="value must be a valid value",
            billable_units=input.value,
        )

    @fal.endpoint("/cuda-exception")
    def cuda_exception(self) -> Output:
        # mimicking error message from PyTorch (https://github.com/pytorch/pytorch/blob/6c65fd03942415b68040e102c44cf5109d2d851e/c10/cuda/CUDACachingAllocator.cpp#L1234C12-L1234C30)
        raise RuntimeError("CUDA out of memory")

    @fal.endpoint("/cuda-exception-2")
    def cuda_exception_2(self) -> Output:
        # https://github.com/pytorch/pytorch/issues/112377
        raise RuntimeError("NVML_SUCCESS == r INTERNAL ASSERT FAILED")

    @fal.endpoint("/cuda-exception-3")
    def cuda_exception_3(self) -> Output:
        raise RuntimeError("cuDNN error: CUDNN_STATUS_INTERNAL_ERROR")


class CancellableApp(fal.App, keep_alive=300, max_concurrency=1, request_timeout=4):
    task = None
    running = 0

    async def _sleep(self, input: Input):
        if self.running > 0:
            raise Exception("App is already running")

        self.task = asyncio.create_task(asyncio.sleep(input.wait_time))
        try:
            self.running += 1
            await self.task
        except asyncio.CancelledError:
            print("Task was cancelled")
            if not self.task.done():
                self.task.cancel()
                with suppress(Exception):
                    await self.task

            raise RequestCancelledException("Request cancelled by the client.")
        finally:
            self.task = None
            self.running -= 1
        return Output(result=input.lhs + input.rhs)

    @fal.endpoint("/")
    async def sleep(self, input: Input) -> Output:
        return await self._sleep(input)

    @fal.endpoint("/well-handled")
    async def well_handled(self, input: Input, request: Request) -> Output:
        async with cancel_on_disconnect(request):
            return await self._sleep(input)

    @fal.endpoint("/cancel")
    async def cancel_handler(self) -> Output:
        if self.task:
            self.task.cancel()
            with suppress(BaseException):
                await self.task
            self.task = None

        return Output(result=0)


class HealthCheckApp(fal.App, keep_alive=300, max_concurrency=1, request_timeout=4):
    @fal.endpoint("/")
    def run(self, input: Input) -> Output:
        time.sleep(input.wait_time)
        return Output(result=input.lhs + input.rhs)

    @fal.endpoint(
        "/health",
        health_check=fal.HealthCheck(
            start_period_seconds=10,
            timeout_seconds=10,
            failure_threshold=3,
            call_regularly=True,
        ),
    )
    def health(self) -> Output:
        return Output(result=0)


class RTInput(BaseModel):
    prompt: str

    def can_batch(
        self,
        other: "RTInput",
        current_batch_size: int = 1,
    ) -> bool:
        return "don't batch" not in other.prompt


class RTOutput(BaseModel):
    text: str


class RTOutputs(BaseModel):
    texts: List[str]


def json_encode_message(message):
    return json.dumps(message, separators=(",", ":")).encode("utf-8")


def json_decode_message(message: bytes):
    return json.loads(message.decode("utf-8"))


class RealtimeApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "S"

    @fal.endpoint("/")
    def generate(self, input: RTInput) -> RTOutput:
        return RTOutput(text=input.prompt)

    @fal.endpoint("/ws", is_websocket=True)
    async def generate_ws(self, websocket: WebSocket) -> None:
        await websocket.accept()
        for _ in range(3):
            await websocket.send_json({"message": "Hello world!"})
        await websocket.close()

    @fal.realtime("/realtime")
    def generate_rt(self, input: RTInput) -> RTOutput:
        return RTOutput(text=input.prompt)

    @fal.realtime("/realtime/server-streaming", buffering=10)
    async def generate_rt_server_streaming(
        self, input: RTInput
    ) -> AsyncIterator[RTOutput]:
        for idx in range(3):
            yield RTOutput(text=f"{input.prompt}:{idx}")

    @fal.realtime("/realtime/server-streaming-sync", buffering=10)
    def generate_rt_server_streaming_sync(self, input: RTInput) -> Iterator[RTOutput]:
        for idx in range(3):
            yield RTOutput(text=f"{input.prompt}:{idx}")

    @fal.realtime("/realtime/client-streaming", session_timeout=0.2)
    async def generate_rt_client_streaming(
        self, inputs: AsyncIterator[RTInput]
    ) -> RTOutputs:
        prompts: List[str] = []
        async for item in inputs:
            prompts.append(item.prompt)
        return RTOutputs(texts=prompts)

    @fal.realtime("/realtime/bidi")
    async def generate_rt_bidi(
        self, inputs: AsyncIterator[RTInput]
    ) -> AsyncIterator[RTOutput]:
        async for item in inputs:
            yield RTOutput(text=f"echo:{item.prompt}")

    @fal.realtime(
        "/realtime/json",
        encode_message=json_encode_message,
        decode_message=json_decode_message,
    )
    def generate_rt_json(self, input: RTInput) -> RTOutput:
        return RTOutput(text=input.prompt)

    @fal.realtime("/realtime/batched", buffering=10, max_batch_size=4)
    def generate_rt_batched(self, input: RTInput, *inputs: RTInput) -> RTOutputs:
        time.sleep(2)  # fixed cost
        return RTOutputs(texts=[input.prompt] + [i.prompt for i in inputs])


class BrokenApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "S"

    @fal.endpoint("/")
    def broken(self) -> Exception:
        raise Exception("this app is designed to fail")


@pytest.fixture(scope="module")
def host() -> Generator[api.Host, None, None]:
    yield addition_app.host


@pytest.fixture(scope="module")
def user(rest_client: Client) -> Generator[User, None, None]:
    user = _get_user(rest_client)
    yield user


@contextmanager
def register_app(
    host: api.FalServerlessHost,
    app: Union[api.ServedIsolatedFunction, api.IsolatedFunction],
    suffix: str = "",
):
    app_alias = str(uuid.uuid4()) + "-test-alias" + ("-" + suffix if suffix else "")
    result = host.register(
        func=app.func,
        options=app.options,
        application_name=app_alias,
        application_auth_mode="private",
        deployment_strategy="recreate",
    )

    assert result
    assert result.result
    assert result.service_urls
    app_revision = result.result.application_id

    try:
        yield app_alias, app_revision
    finally:
        with host._connection as client:
            client.delete_alias(app_alias)


@pytest.fixture(scope="module")
def base_app(host: api.FalServerlessHost):
    # running apps without aliases is no longer supported
    # so we need to create an alias for the app
    with register_app(host, addition_app, "base") as (app_alias, app_revision):
        yield app_alias, app_revision


@pytest.fixture(scope="module")
def test_app(host: api.FalServerlessHost, user: User):
    with register_app(host, addition_app, "addition") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_nomad_app(host: api.FalServerlessHost, user: User):
    with register_app(host, nomad_addition_app, "nomad") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_container_app(host: api.FalServerlessHost, user: User):
    with register_app(host, container_addition_app, "container") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_container_build_args_app(host: api.FalServerlessHost, user: User):
    with register_app(host, container_build_args_app, "build-args") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_fastapi_app(host: api.FalServerlessHost, user: User):
    with register_app(host, calculator_app, "fastapi") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_stateful_app(host: api.FalServerlessHost, user: User):
    stateful_app = wrap_app(StatefulAdditionApp)
    with register_app(host, stateful_app, "stateful") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_pydantic_validation_error():
    with AppClient.connect(StatefulAdditionApp) as client:
        yield client


@pytest.fixture(scope="module")
def test_cancellable_app(host: api.FalServerlessHost, user: User):
    cancellable_app = wrap_app(CancellableApp)
    with register_app(host, cancellable_app, "cancellable") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_exception_app():
    with AppClient.connect(ExceptionApp) as client:
        yield client


@pytest.fixture(scope="module")
def test_health_check_app(host: api.FalServerlessHost, user: User):
    health_check_app = wrap_app(HealthCheckApp)
    with register_app(host, health_check_app, "health-check") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_sleep_app(host: api.FalServerlessHost, user: User):
    sleep_app = wrap_app(SleepApp)
    with register_app(host, sleep_app, "sleep") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_queue_blocking_app(host: api.FalServerlessHost, user: User):
    queue_blocking_app = wrap_app(QueueBlockingApp)
    with register_app(host, queue_blocking_app, "queue-blocking") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


@pytest.fixture(scope="module")
def test_realtime_app(host: api.FalServerlessHost, user: User):
    realtime_app = wrap_app(RealtimeApp)
    with register_app(host, realtime_app, "realtime") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


def test_broken_app_failure(host: api.FalServerlessHost, user: User):
    with pytest.raises(FalServerlessException) as e:
        wrap_app(BrokenApp)

    assert "Failed to generate OpenAPI" in str(e)


@pytest.mark.flaky(max_runs=3)
def test_app_client(test_app: str, test_nomad_app: str):
    response = apps.run(test_app, arguments={"lhs": 1, "rhs": 2})
    assert response["result"] == 3

    response = apps.run(test_app, arguments={"lhs": 2, "rhs": 3, "wait_time": 1})
    assert response["result"] == 5

    response = apps.run(test_nomad_app, arguments={"lhs": 1, "rhs": 2})
    assert response["result"] == 3

    response = apps.run(test_nomad_app, arguments={"lhs": 2, "rhs": 3, "wait_time": 1})
    assert response["result"] == 5


def test_ws_client(test_app: str):
    with apps.ws(test_app) as connection:
        for i in range(3):
            response = json.loads(connection.run({"lhs": 1, "rhs": i}))
            assert response["result"] == 1 + i

        for i in range(3):
            connection.send({"lhs": 2, "rhs": i})

        for i in range(3):
            # they should be in order
            response = json.loads(connection.recv())
            assert response["result"] == 2 + i


def test_app_client_old_format(base_app: Tuple[str, str], user: User):
    app_alias, _ = base_app
    old_format = f"{user.user_id}-{app_alias}"
    response = apps.run(old_format, arguments={"lhs": 1, "rhs": 2})
    assert response["result"] == 3


def test_app_client_path_included_in_app_id(test_stateful_app: str):
    response = apps.run(test_stateful_app + "/reset", arguments={})
    assert response["result"] == 0

    response = apps.run(test_stateful_app + "/increment", arguments={"value": 3})
    assert response["result"] == 3

    # if put in path we do not need to prefix with /
    response = apps.run(test_stateful_app, arguments={"value": 3}, path="increment")
    assert response["result"] == 6


def test_stateful_app_client(test_stateful_app: str):
    response = apps.run(test_stateful_app, arguments={}, path="/reset")
    assert response["result"] == 0

    response = apps.run(test_stateful_app, arguments={"value": 1}, path="/increment")
    assert response["result"] == 1

    response = apps.run(test_stateful_app, arguments={"value": 2}, path="/increment")
    assert response["result"] == 3

    response = apps.run(test_stateful_app, arguments={"value": 1}, path="/decrement")
    assert response["result"] == 2

    response = apps.run(test_stateful_app, arguments={"value": 2}, path="/decrement")
    assert response["result"] == 0


@pytest.mark.flaky(max_runs=3)
def test_app_cancellation(test_app: str, test_cancellable_app: str):
    request_handle = apps.submit(
        test_cancellable_app, arguments={"lhs": 1, "rhs": 2, "wait_time": 6}
    )

    while True:
        status = request_handle.status()
        time.sleep(0.05)
        if isinstance(status, apps.InProgress):
            # The app is running
            break

    # cancel the request
    request_handle.cancel()

    # should still finish successfully and return 499
    with pytest.raises(HTTPStatusError) as e:
        request_handle.get()
    assert e.value.response.status_code == 499

    # normal app should just ignore the cancellation
    request_handle = apps.submit(
        test_app, arguments={"lhs": 1, "rhs": 2, "wait_time": 6}
    )

    while True:
        status = request_handle.status()
        time.sleep(0.05)
        if isinstance(status, apps.InProgress):
            # The app is running
            break

    # cancel the request
    request_handle.cancel()

    response = request_handle.get()
    assert response == {"result": 3}


@pytest.mark.flaky(max_runs=3)
def test_app_disconnect_behavior(test_app: str, test_cancellable_app: str):
    with pytest.raises(HTTPStatusError) as e:
        apps.run(
            test_cancellable_app,
            arguments={"lhs": 1, "rhs": 2, "wait_time": 6},
            path="/well-handled",
        )
    assert (
        e.value.response.status_code == 504
    ), "Expected Gateway Timeout even though the app handled it"

    # and running it again shows the app "handled" it
    response = apps.run(
        test_cancellable_app,
        arguments={"lhs": 1, "rhs": 2, "wait_time": 1},
        path="/well-handled",
    )
    assert response == {"result": 3}

    # vs on an unhandled one

    with pytest.raises(HTTPStatusError) as e:
        apps.run(
            test_cancellable_app,
            arguments={"lhs": 1, "rhs": 2, "wait_time": 6},
        )
    assert (
        e.value.response.status_code == 504
    ), "Expected Gateway Timeout even though the app handled it"


@pytest.mark.flaky(max_runs=3)
def test_start_timeout_queue_blocking(test_queue_blocking_app: str):
    """
    Test that start_timeout correctly times out a request waiting in queue.

    Scenario:
    1. Send a 10-second sleep request (occupies the only slot)
    2. While it's processing, send a second request with start_timeout=5
    3. The second request should return 504 because it times out waiting in queue
       (before processing starts)
    4. First request should complete successfully
    """
    import fal_client
    from fal_client.client import FalClientHTTPError

    # Send a long-running request that will occupy the only slot
    # (max_concurrency=1, max_multiplexing=1)
    # Use 10 seconds to ensure it blocks long enough for the second request to timeout
    first_handle = apps.submit(test_queue_blocking_app, arguments={"wait_time": 10})

    # Wait for the first request to start processing
    while True:
        status = first_handle.status()
        if isinstance(status, apps.InProgress):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(0.1)
        else:
            raise Exception(f"Unexpected status for first request: {status}")

    # Now send a second request with a short start_timeout
    # This should fail because it will timeout waiting in the queue
    with pytest.raises(FalClientHTTPError) as exc_info:
        fal_client.subscribe(
            test_queue_blocking_app,
            arguments={"wait_time": 1},
            start_timeout=5,
        )

    # Should get a 504 timeout error
    assert (
        exc_info.value.status_code == 504
    ), f"Expected 504 timeout, got {exc_info.value.status_code}"

    # Verify the timeout type header indicates it was a user timeout
    timeout_type = exc_info.value.response_headers.get("x-fal-request-timeout-type")
    assert timeout_type == "user", f"Expected 'user' timeout type, got {timeout_type}"

    # First request should complete successfully
    result = first_handle.get()
    assert result == {"slept": True}, f"First request should succeed, got {result}"


@pytest.mark.xfail(
    reason="Temporary disabled while investigating backend issue. Ping @efiop"
)
@pytest.mark.flaky(max_runs=3)
def test_app_client_async(test_sleep_app: str):
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 1})
    with pytest.raises(HTTPStatusError) as e:
        # Not yet completed
        handle.fetch_result()

    assert e.value.response.status_code == 400

    # Wait until the app is completed
    assert handle.get() == {"slept": True}

    # New request
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 5})

    for event in handle.iter_events(logs=True):
        assert isinstance(event, (apps.Queued, apps.InProgress))
        if isinstance(event, apps.InProgress) and event.logs:
            logs = [log["message"] for log in event.logs]
            assert "sleeping..." in logs
        elif isinstance(event, apps.Queued):
            assert event.position == 0

    for _ in range(10):
        status = handle.status(logs=True)
        assert isinstance(status, apps.Completed)
        if status.logs:
            break

    assert status.logs, "Logs missing from Completed status"
    assert any("sleeping..." in log["message"] for log in status.logs)

    # It is safe to use fetch_result when we know for a fact the request itself
    # is completed.
    result = handle.fetch_result()

    # .get() can still be used and will return the same value
    get_result = handle.get()
    assert result == get_result
    assert result == {"slept": True}


# If the logging subsystem is not working for some nodes, this test will flake
@pytest.mark.xfail(
    reason="Temporary disabled while investigating backend issue. Ping @efiop"
)
def test_traceback_logs(test_exception_app: AppClient, rest_client: Client):
    date = (
        datetime.now(timezone.utc).replace(tzinfo=None) - timedelta(seconds=1)
    ).isoformat()

    with pytest.raises(AppClientError):
        test_exception_app.fail({})

    with httpx.Client(
        base_url=rest_client.base_url,
        headers=rest_client.get_headers(),
        timeout=300,
    ) as client:
        # Give some time for logs to propagate through the logging subsystem.
        for _ in range(10):
            time.sleep(2)
            response = client.get(
                rest_client.base_url + f"/logs/?traceback=true&since={date}"
            )

            logs = response.json()
            if len(logs) > 0:
                break

        assert len(logs) > 0
        for log in logs:
            assert log["message"].count("\n") > 1, "Logs should be multi-line"
            assert (
                '{"traceback":' not in log["message"]
            ), "Logs should not be JSON-wrapped"
            assert (
                "this app is designed to fail" in log["message"]
            ), "Logs should contain the traceback message"


def test_app_openapi_spec_metadata(
    base_app: Tuple[str, str], user: User, rest_client: Client
):
    app_alias, _ = base_app
    res = app_metadata.sync_detailed(
        app_alias_or_id=app_alias, app_user_id=user.user_id, client=rest_client
    )

    assert res.status_code == 200, f"Failed to fetch metadata for app {app_alias}"
    assert res.parsed, f"Failed to parse metadata for app {app_alias}"

    metadata = res.parsed.to_dict()
    assert "openapi" in metadata, f"openapi key missing from metadata {metadata}"
    openapi_spec: dict = metadata["openapi"]
    for key in ["openapi", "info", "paths", "components"]:
        assert key in openapi_spec, f"{key} key missing from openapi {openapi_spec}"


def test_app_no_serve_spec_metadata(test_fastapi_app: str, rest_client: Client):
    # We do not store the openapi spec for apps that do not use serve=True
    user_id, _, app_id = test_fastapi_app.partition("/")
    res = app_metadata.sync_detailed(
        app_alias_or_id=app_id, app_user_id=user_id, client=rest_client
    )

    assert (
        res.status_code == 200
    ), f"Failed to fetch metadata for app {test_fastapi_app}"
    assert res.parsed, f"Failed to parse metadata for app {test_fastapi_app}"

    metadata = res.parsed.to_dict()
    assert (
        "openapi" not in metadata
    ), f"openapi should not be present in metadata {metadata}"


def test_404_response(test_app: str, request: pytest.FixtureRequest):
    with pytest.raises(HTTPStatusError, match="Path /.*other not found"):
        apps.run(test_app, path="/other", arguments={"lhs": 1, "rhs": 2})


def test_404_billable_units(test_exception_app: AppClient):
    """Test that 404 responses include x-fal-billable-units: 0 header."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/non-existent-endpoint"
        response = httpx_client.post(
            url,
            json={},
            timeout=30,
        )

        assert response.status_code == 404
        assert response.headers.get("x-fal-billable-units") == "0"


def test_app_no_auth():
    # This will just pass for users with shared apps access
    app_alias = str(uuid.uuid4()) + "-alias"
    with pytest.raises(api.FalServerlessError, match="Must specify auth_mode"):
        addition_app.host.register(
            func=addition_app.func,
            options=addition_app.options,
            # random enough
            application_name=app_alias,
            deployment_strategy="recreate",
        )


def test_app_deploy_scale(host: api.FalServerlessHost):
    from dataclasses import replace

    app_alias = str(uuid.uuid4()) + "-alias"
    result = addition_app.host.register(
        func=addition_app.func,
        options=addition_app.options,
        application_name=app_alias,
        application_auth_mode="private",
        deployment_strategy="recreate",
    )
    assert result
    assert result.result
    assert result.service_urls
    app_revision = result.result.application_id

    options = replace(
        addition_app.options,
        host={
            **addition_app.options.host,
            "max_multiplexing": 3,
            "max_concurrency": 2,
        },
    )
    kwargs = dict(
        func=addition_app.func,
        options=options,
        application_name=app_alias,
        application_auth_mode="private",
        deployment_strategy="recreate",
    )

    result = addition_app.host.register(**kwargs, scale=False)
    assert result
    assert result.result
    assert result.service_urls
    app_revision = result.result.application_id

    with host._connection as client:
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert found, f"Could not find app {app_alias} in {res}"
        assert found.revision == app_revision
        # multiplexing is revision-specific
        assert found.max_multiplexing == 3, "Expected max_multiplexing to have changed"
        # max_concurrency is alias-specific
        assert found.max_concurrency == 1, "Expected max_concurrency to stay the same"

    result = addition_app.host.register(**kwargs, scale=True)
    assert result
    assert result.result
    assert result.service_urls
    app_revision = result.result.application_id

    with host._connection as client:
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert found, f"Could not find app {app_alias} in {res}"
        assert found.revision == app_revision
        # when scaling, all values are updated
        assert found.max_multiplexing == 3
        assert found.max_concurrency == 2


# List aliases is taking long
@pytest.mark.timeout(600)
def test_app_update_app(base_app: Tuple[str, str]):
    app_alias, app_revision = base_app

    host: api.FalServerlessHost = addition_app.host  # type: ignore
    with host._connection as client:
        # Get the registered values
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert found, f"Could not find app {app_alias} in {res}"
        assert found.revision == app_revision

    with host._connection as client:
        new_keep_alive = found.keep_alive + 1
        new_max_concurrency = found.max_concurrency + 1
        new_max_multiplexing = found.max_multiplexing + 1

        res = client.update_application(
            application_name=app_alias,
            keep_alive=new_keep_alive,
            max_concurrency=new_max_concurrency,
            max_multiplexing=new_max_multiplexing,
        )
        assert res.alias == app_alias
        assert res.keep_alive == new_keep_alive
        assert res.max_concurrency == new_max_concurrency
        assert res.max_multiplexing == new_max_multiplexing

    with host._connection as client:
        new_keep_alive = new_keep_alive + 1
        res = client.update_application(
            application_name=app_alias,
            keep_alive=new_keep_alive,
        )
        assert res.alias == app_alias
        assert res.keep_alive == new_keep_alive
        assert res.max_concurrency == new_max_concurrency
        assert res.max_multiplexing == new_max_multiplexing

    with host._connection as client:
        new_max_concurrency = new_max_concurrency - 1
        res = client.update_application(
            application_name=app_alias,
            max_concurrency=new_max_concurrency,
        )
        assert res.alias == app_alias
        assert res.keep_alive == new_keep_alive
        assert res.max_concurrency == new_max_concurrency
        assert res.max_multiplexing == new_max_multiplexing


# List aliases is taking long
@pytest.mark.timeout(600)
def test_app_set_delete_alias(base_app: Tuple[str, str]):
    app_alias, app_revision = base_app

    host: api.FalServerlessHost = addition_app.host  # type: ignore

    with host._connection as client:
        # Get the registered values
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert found, f"Could not find app {app_alias} in {res}"
        assert found.revision == app_revision
        assert found.auth_mode == "private"

    with host._connection as client:
        # Get the registered values
        res = client.create_alias(app_alias, app_revision, "public")

    with host._connection as client:
        # Get the registered values
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert found, f"Could not find app {app_alias} in {res}"
        assert found.revision == app_revision
        assert found.auth_mode == "public"

    with host._connection as client:
        res = client.delete_alias(alias=app_alias)
        assert res == app_revision

    with host._connection as client:
        # Get the registered values
        res = client.list_aliases()
        found = next(filter(lambda alias: alias.alias == app_alias, res), None)
        assert not found, f"Found app {app_alias} in {res} after deletion"


@pytest.mark.flaky(max_runs=3)
def test_realtime_connection(test_realtime_app):
    response = apps.run(test_realtime_app, arguments={"prompt": "a cat"})
    assert response["text"] == "a cat"

    with apps._connect(test_realtime_app) as connection:
        for _ in range(3):
            response = connection.run({"prompt": "a cat"})
            assert response["text"] == "a cat"

    with apps._connect(test_realtime_app, path="/realtime/batched") as connection:
        connection.send({"prompt": "keep busy"})
        time.sleep(1)

        for prompt in range(10):
            connection.send({"prompt": str(prompt)})

        assert connection.recv()["texts"] == ["keep busy"]

        received_prompts = set()
        batch_sizes = []
        while len(received_prompts) < 10:
            response = connection.recv()
            received_prompts.update(response["texts"])
            batch_sizes.append(len(response["texts"]))

        assert len(received_prompts) == 10
        assert batch_sizes == [4, 4, 2]


def test_realtime_ws_endpoint(test_realtime_app):
    app_id = apps._backwards_compatible_app_id(test_realtime_app)
    url = apps._REALTIME_URL_FORMAT.format(app_id=app_id) + "/ws"
    creds = get_default_credentials()

    with ws_client.connect(
        url, additional_headers=creds.to_headers(), open_timeout=90
    ) as ws:
        messages = []
        for _ in range(3):
            payload = ws.recv()
            if isinstance(payload, bytes):
                payload = payload.decode("utf-8")
            messages.append(json.loads(payload))

    assert messages == [{"message": "Hello world!"}] * 3


def test_realtime_connection_custom_codec(test_realtime_app):
    with apps._connect(
        test_realtime_app,
        path="/realtime/json",
        encode_message=json_encode_message,
        decode_message=json_decode_message,
    ) as connection:
        response = connection.run({"prompt": "json cat"})
        assert response["text"] == "json cat"


def test_realtime_server_streaming_mode(test_realtime_app):
    with apps._connect(
        test_realtime_app, path="/realtime/server-streaming"
    ) as connection:
        connection.send({"prompt": "stream"})
        responses = [connection.recv() for _ in range(3)]
        assert [response["text"] for response in responses] == [
            "stream:0",
            "stream:1",
            "stream:2",
        ]


def test_realtime_server_streaming_sync_mode(test_realtime_app):
    with apps._connect(
        test_realtime_app, path="/realtime/server-streaming-sync"
    ) as connection:
        connection.send({"prompt": "stream"})
        responses = [connection.recv() for _ in range(3)]
        assert [response["text"] for response in responses] == [
            "stream:0",
            "stream:1",
            "stream:2",
        ]


def test_realtime_client_streaming_mode(test_realtime_app):
    with apps._connect(
        test_realtime_app, path="/realtime/client-streaming"
    ) as connection:
        connection.send({"prompt": "first"})
        connection.send({"prompt": "second"})
        connection.send({"prompt": "third"})
        response = connection.recv()
        assert response["texts"] == ["first", "second", "third"]


def test_realtime_bidi_mode(test_realtime_app):
    with apps._connect(test_realtime_app, path="/realtime/bidi") as connection:
        connection.send({"prompt": "one"})
        connection.send({"prompt": "two"})
        assert connection.recv()["text"] == "echo:one"
        assert connection.recv()["text"] == "echo:two"


@contextmanager
def delete_workflow_on_exit(client: httpx.Client, workflow_url: str):
    try:
        yield
    finally:
        client.delete(workflow_url)


def test_workflows(test_app: str, rest_client: Client):
    workflow = Workflow(
        name="test_workflow_" + secrets.token_hex(),
        input_schema={},
        output_schema={},
    )
    # (lhs + rhs) + (lhs + rhs)
    lhs = workflow.run(
        test_app,
        input={
            "lhs": workflow.input.lhs,
            "rhs": workflow.input.rhs,
        },
    )
    rhs = workflow.run(
        test_app,
        input={
            "lhs": workflow.input.lhs,
            "rhs": workflow.input.rhs,
        },
    )
    out = workflow.run(
        test_app,
        input={
            "lhs": lhs.result,
            "rhs": rhs.result,
        },
    )
    workflow.set_output({"result": out.result})
    workflow_id = workflow.publish(title="Test Workflow", is_public=False)

    # Test the underlying app
    data = fal.apps.run(test_app, arguments={"lhs": 2, "rhs": 3})
    assert data["result"] == 5

    with httpx.Client(
        base_url=rest_client.base_url,
        headers=rest_client.get_headers(),
        timeout=300,
    ) as client:
        with delete_workflow_on_exit(
            client, rest_client.base_url + "/workflows/" + workflow_id
        ):
            data = fal.apps.run(
                "workflows/" + workflow_id, arguments={"lhs": 2, "rhs": 3}
            )
            assert data["result"] == 10


def test_app_exceptions(test_exception_app: AppClient):
    with pytest.raises(AppClientError) as app_exc:
        test_exception_app.app_exception({})

    assert app_exc.value.status_code == 401

    with pytest.raises(AppClientError) as field_exc:
        test_exception_app.field_exception({"lhs": 1, "rhs": "2"})

    assert field_exc.value.status_code == 422

    assert field_exc.value.headers.get("x-fal-billable-units") is None

    with pytest.raises(AppClientError) as cuda_exc:
        test_exception_app.cuda_exception({})

    assert cuda_exc.value.status_code == _CUDA_OOM_STATUS_CODE
    assert _CUDA_OOM_MESSAGE in cuda_exc.value.message

    with pytest.raises(AppClientError) as cuda_exc:
        test_exception_app.cuda_exception_2({})

    assert cuda_exc.value.status_code == _CUDA_OOM_STATUS_CODE
    assert _CUDA_OOM_MESSAGE in cuda_exc.value.message

    with pytest.raises(AppClientError) as cuda_exc:
        test_exception_app.cuda_exception_3({})

    assert cuda_exc.value.status_code == _CUDA_OOM_STATUS_CODE
    assert _CUDA_OOM_MESSAGE in cuda_exc.value.message


def test_pydantic_validation_billing(test_pydantic_validation_error: AppClient):
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_pydantic_validation_error.url + "/increment"
        response = httpx_client.post(
            url,
            json={"value": "this-is-not-an-integer"},
            timeout=30,
        )

        assert response.status_code == 422
        assert response.headers.get("x-fal-billable-units") == "0"


def test_field_exception_billing(test_exception_app: AppClient):
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception"
        response = httpx_client.post(
            url,
            json={"lhs": 1, "rhs": 2},
            timeout=30,
        )

        assert response.status_code == 422
        # For errors raised on runtime, developers should be handling the billing.
        # Therefore not adding billing units.
        assert not hasattr(response.headers, "x-fal-billable-units")


def test_field_exception_int_billable_units_formatting(test_exception_app: AppClient):
    """Test that int billable_units are formatted without decimal places."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception-units"
        response = httpx_client.post(
            url,
            json={"value": 42},
            timeout=30,
        )

        assert response.status_code == 422
        assert response.headers.get("x-fal-billable-units") == "42"


def test_field_exception_float_billable_units_formatting(test_exception_app: AppClient):
    """Test that float billable_units are formatted with 8 decimal places."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception-units"
        response = httpx_client.post(
            url,
            json={"value": 3.14159265},
            timeout=30,
        )

        assert response.status_code == 422
        assert response.headers.get("x-fal-billable-units") == "3.14159265"


def test_field_exception_scientific_notation_small(test_exception_app: AppClient):
    """Test that small scientific notation values are properly formatted."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception-units"
        response = httpx_client.post(
            url,
            json={"value": 1.23e-5},
            timeout=30,
        )

        assert response.status_code == 422
        # 1.23e-5 = 0.0000123 (float type uses .8f format)
        assert response.headers.get("x-fal-billable-units") == "0.00001230"


def test_field_exception_scientific_notation_large(test_exception_app: AppClient):
    """Test that large scientific notation values are properly formatted."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception-units"
        response = httpx_client.post(
            url,
            json={"value": 1.23e10},
            timeout=30,
        )

        assert response.status_code == 422
        # 1.23e10 = 12300000000.0 (float type uses .8f format)
        assert response.headers.get("x-fal-billable-units") == "12300000000.00000000"


def test_field_exception_invalid_billable_units(test_exception_app: AppClient):
    """Test that invalid billable_units (non-numeric string) raises an error."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception-units"
        response = httpx_client.post(
            url,
            json={"value": "not_a_number"},
            timeout=30,
        )

        # should return 500 internal server error due to ValueError when
        # converting to float
        assert response.status_code == 500


def test_field_exception_default_billable_units(test_exception_app: AppClient):
    """Test that when billable_units is not set (None), no header is included."""
    with httpx.Client(headers=_auth_headers()) as httpx_client:
        url = test_exception_app.url + "/field-exception"
        response = httpx_client.post(
            url,
            json={"lhs": 1, "rhs": 2},
            timeout=30,
        )

        assert response.status_code == 422
        # When billable_units is None (default), header should not be present
        assert "x-fal-billable-units" not in response.headers


def submit_and_wait_for_runner(app: str, arguments: dict = {}, *, path: str = ""):
    handle = apps.submit(app, arguments=arguments, path=path)

    while True:
        status = handle.status()
        if isinstance(status, apps.InProgress) or isinstance(status, apps.Completed):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(0.1)
        else:
            raise Exception(f"Failed to start the app: {status}")

    return handle


@pytest.mark.flaky(max_runs=3)
def test_stop_runner(host: api.FalServerlessHost, test_sleep_app: str):
    # Submit a runner and wait for it to be idle
    submit_and_wait_for_runner(test_sleep_app, arguments={"wait_time": 1})

    with host._connection as client:
        timeout = 10
        start_time = time.time()
        while True:
            _, _, app_alias = test_sleep_app.partition("/")
            runners = client.list_alias_runners(app_alias)
            assert len(runners) == 1

            if runners[0].in_flight_requests == 0:
                break
            elif time.time() - start_time > timeout:
                raise Exception(f"Timeout waiting for runner to be idle: {runners[0]}")
            time.sleep(1)

    # Because the runner is not requested to be stopped, it should be reused
    submit_and_wait_for_runner(test_sleep_app, arguments={"wait_time": 1})

    with host._connection as client:
        _, _, app_alias = test_sleep_app.partition("/")
        runners = client.list_alias_runners(app_alias)
        assert len(runners) == 1

    # Request to stop the runner
    with host._connection as client:
        with pytest.raises(Exception) as e:
            client.stop_runner("1234567890")

        assert "not found" in str(e).lower()

        _, _, app_alias = test_sleep_app.partition("/")
        runners = client.list_alias_runners(app_alias)
        assert len(runners) == 1

        client.stop_runner(runners[0].runner_id)

    # Because the runner is requested to be stopped,
    # it should not be reused and a new runner should be created
    submit_and_wait_for_runner(test_sleep_app, arguments={"wait_time": 1})

    with host._connection as client:
        runners = client.list_alias_runners(app_alias)
        assert len(runners) == 2


@pytest.mark.flaky(max_runs=3)
def test_kill_runner(host: api.FalServerlessHost, test_sleep_app: str):
    # Kill all the replicas of the app that is already running
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 10})

    while True:
        status = handle.status()
        if isinstance(status, apps.InProgress):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(1)
        else:
            raise Exception(f"Failed to start the app: {status}")

    with host._connection as client:
        with pytest.raises(Exception) as e:
            client.kill_runner("1234567890")

        assert "not found" in str(e).lower()

        _, _, app_alias = test_sleep_app.partition("/")
        runners = client.list_alias_runners(app_alias)
        existing_runners = len(
            [runner for runner in runners if runner.state == RunnerState.RUNNING]
        )

        client.kill_runner(runners[0].runner_id)

        runners = client.list_alias_runners(app_alias)
        num_runners = len(
            [runner for runner in runners if runner.state == RunnerState.RUNNING]
        )
        assert num_runners == existing_runners - 1


@pytest.mark.flaky(max_runs=3)
def test_rollout_application(host: api.FalServerlessHost, test_sleep_app: str):
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 30})

    while True:
        status = handle.status()
        if isinstance(status, apps.InProgress):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(1)
        else:
            raise Exception(f"Failed to start the app: {status}")

    with host._connection as client:
        _, _, app_alias = test_sleep_app.partition("/")
        runners_before = client.list_alias_runners(app_alias)
        assert len(runners_before) == 1
        runner_id_before = runners_before[0].runner_id

        client.rollout_application(app_alias, force=True)

        time.sleep(15)

        runners_after = client.list_alias_runners(app_alias)
        runner_ids_after = {r.runner_id for r in runners_after}

        assert runner_id_before not in runner_ids_after

        client.rollout_application(app_alias, force=True)

        time.sleep(3)

        runners_final = client.list_alias_runners(app_alias)
        runner_ids_final = {r.runner_id for r in runners_final}

        assert not runner_ids_after.intersection(runner_ids_final)


@pytest.mark.flaky(max_runs=3)
def test_shell_runner(host: api.FalServerlessHost, test_sleep_app: str):
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 30})

    while True:
        status = handle.status()
        if isinstance(status, apps.InProgress):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(1)
        else:
            raise Exception(f"Failed to start the app: {status}")

    with host._connection as client:
        _, _, app_alias = test_sleep_app.partition("/")
        runners = client.list_alias_runners(app_alias)
        assert len(runners) == 1
        runner_id = runners[0].runner_id

        proc = subprocess.Popen(
            ["python", "-m", "fal", "runners", "shell", runner_id],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        try:
            commands = b"echo 'a' > t.txt\ncat t.txt\nexit\n"
            stdout, stderr = proc.communicate(input=commands, timeout=10)
            assert b"a" in stdout, f"Expected 'a' in output, got: {stdout.decode()}"
        finally:
            if proc.poll() is None:
                proc.kill()
                proc.wait()


@pytest.mark.flaky(max_runs=3)
def test_exec_runner(host: api.FalServerlessHost, test_sleep_app: str):
    handle = apps.submit(test_sleep_app, arguments={"wait_time": 30})

    while True:
        status = handle.status()
        if isinstance(status, apps.InProgress):
            break
        elif isinstance(status, apps.Queued):
            time.sleep(1)
        else:
            raise Exception(f"Failed to start the app: {status}")

    with host._connection as client:
        _, _, app_alias = test_sleep_app.partition("/")
        runners = client.list_alias_runners(app_alias)
        assert len(runners) == 1
        runner_id = runners[0].runner_id

        proc = subprocess.Popen(
            [
                "python",
                "-m",
                "fal",
                "runners",
                "exec",
                runner_id,
                "--",
                "echo",
                "hello",
            ],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        try:
            stdout, stderr = proc.communicate(timeout=10)
            assert (
                b"hello" in stdout
            ), f"Expected 'hello' in output, got: {stdout.decode()}"
        finally:
            if proc.poll() is None:
                proc.kill()
                proc.wait()


def test_container_app_client(test_container_app: str):
    response = apps.run(test_container_app, arguments={"lhs": 1, "rhs": 2})
    assert response["result"] == 3


def test_container_build_args_app_client(test_container_build_args_app: str):
    response = apps.run(test_container_build_args_app, {})
    assert response == "built with build args"


def test_container_no_cache_app_client(host: api.FalServerlessHost):
    stdout = StringIO()
    with redirect_stdout(stdout):
        with register_app(host, container_no_cache_app, "container") as (app_alias, _):
            pass

    stdout = StringIO()
    with redirect_stdout(stdout):
        with register_app(host, container_cache_enabled_app, "container") as (
            app_alias,
            _,
        ):
            pass

    logs = stdout.getvalue()
    assert re.search(r"Image \S+ already exists", logs) is not None

    stdout = StringIO()
    with redirect_stdout(stdout):
        with register_app(host, container_no_cache_app, "container") as (app_alias, _):
            pass

    logs = stdout.getvalue()
    assert re.search(r"Image \S+ already exists", logs) is None


class HintsApp(fal.App, keep_alive=300, max_concurrency=1):
    machine_type = "S"

    def provide_hints(self) -> List[str]:
        return ["é", "😀"]

    @fal.endpoint("/add")
    def add(self, input: Input) -> Output:
        return Output(result=input.lhs + input.rhs)


def test_hints_encoding():
    """
    Make sure that hints that can't be encoded in latin-1 don't crash the app
    https://github.com/encode/starlette/blob/a766a58d14007f07c0b5782fa78cdc370b892796/starlette/datastructures.py#L568
    """
    with AppClient.connect(HintsApp) as client:
        with httpx.Client(headers=_auth_headers()) as httpx_client:
            url = client.url + "/add"
            resp = httpx_client.post(
                url,
                json={"lhs": 1, "rhs": 2},
                timeout=30,
            )
            assert resp.is_success
            assert resp.json()["result"] == 3


def _external_get_request_id() -> str:
    request_id = get_current_app().current_request.headers.get("x-request-id", "")
    return request_id


class AppRefOutput(BaseModel):
    from_app: str
    from_external_method: str


class AppRefApp(fal.App, keep_alive=300, max_concurrency=1, max_multiplexing=3):
    machine_type = "XS"

    @fal.endpoint("/")
    def run(self, request: Request) -> AppRefOutput:
        request_id = request.headers.get("x-request-id", "")

        # sleep to intentionally cause a race condition
        time.sleep(3)

        return AppRefOutput(
            from_app=request_id,
            from_external_method=_external_get_request_id(),
        )


@pytest.fixture(scope="module")
def test_app_ref_app(host: api.FalServerlessHost, user: User):
    app_ref_app = wrap_app(AppRefApp)
    with register_app(host, app_ref_app, "app-ref") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


def test_app_ref_app_client(test_app_ref_app: str):
    time.sleep(3)

    handle_1 = apps.submit(test_app_ref_app, arguments={})
    time.sleep(1)
    handle_2 = apps.submit(test_app_ref_app, arguments={})
    time.sleep(1)
    handle_3 = apps.submit(test_app_ref_app, arguments={})
    time.sleep(1)

    result_1 = handle_1.get()
    result_2 = handle_2.get()
    result_3 = handle_3.get()

    assert result_1["from_app"] == result_1["from_external_method"]
    assert result_2["from_app"] == result_2["from_external_method"]
    assert result_3["from_app"] == result_3["from_external_method"]


class SetUUIDInput(BaseModel):
    uuid: str


# for now it only works in newly built containers
class GracefulShutdownApp(
    fal.App,
    keep_alive=300,
    max_concurrency=1,
    image=fal.ContainerImage.from_dockerfile_str(
        f"""
FROM python:{actual_python}-slim
RUN apt-get update \
    && apt-get install -y --no-install-recommends git curl \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*
"""
    ),
    skip_retry_conditions=["server_error", "connection_error", "timeout"],
):
    machine_type = "XS"
    latest_request_id = None
    uuid = None

    def setup(self):
        self.stop = False

    def handle_exit(self):
        self.stop = True

    @fal.endpoint("/set-uuid")
    async def set_uuid(self, input: SetUUIDInput) -> str:
        self.uuid = input.uuid
        return "ok"

    @fal.endpoint("/")
    async def sleep(self, input: SleepInput, request: Request) -> SleepOutput:
        request_id = request.headers.get(REQUEST_ID_KEY)
        if request_id is None:
            raise Exception("No request id found")

        self.latest_request_id = request_id

        for i in range(input.wait_time):
            print(f"sleeping {i + 1} of {input.wait_time}...", flush=True)
            await asyncio.sleep(1)
        return SleepOutput(slept=True)

    @fal.endpoint("/with-stop")
    async def sleep_with_stop(self, input: SleepInput, request: Request) -> SleepOutput:
        request_id = request.headers.get(REQUEST_ID_KEY)
        if request_id is None:
            raise Exception("No request id found")

        self.latest_request_id = request_id

        for i in range(input.wait_time):
            if self.stop:
                break
            print(f"sleeping {i + 1} of {input.wait_time}...", flush=True)
            await asyncio.sleep(1)

        return SleepOutput(slept=True)

    @fal.endpoint("/latest-request-id")
    async def fetch_latest_request_id(self) -> str:
        if self.uuid is None:
            return ""

        try:
            with open(f"/data/teardown/{self.uuid}.txt") as f:
                latest_request_id = f.read()

            os.unlink(f"/data/teardown/{self.uuid}.txt")
            return latest_request_id
        except Exception as e:
            return str(e)

    def teardown(self):
        if self.latest_request_id is None or self.uuid is None:
            return

        os.makedirs("/data/teardown", exist_ok=True)
        with open(f"/data/teardown/{self.uuid}.txt", "w") as f:
            f.write(self.latest_request_id)


@pytest.fixture(scope="module")
def test_graceful_shutdown_app(host: api.FalServerlessHost, user: User):
    graceful_shutdown_app = wrap_app(GracefulShutdownApp)
    with register_app(host, graceful_shutdown_app, "graceful-shutdown") as (
        app_alias,
        _,
    ):
        yield f"{user.username}/{app_alias}"


def graceful_shutdown(
    test_graceful_shutdown_app: str,
    host: api.FalServerlessHost,
    *,
    wait_time: int,
    path: str,
) -> bool:
    time.sleep(2)

    token = str(uuid.uuid4())
    assert (
        apps.run(test_graceful_shutdown_app, {"uuid": token}, path="/set-uuid") == "ok"
    ), "UUID not set"

    handle = submit_and_wait_for_runner(
        test_graceful_shutdown_app, arguments={"wait_time": wait_time}, path=path
    )
    saved_request_id = handle.request_id

    time.sleep(2)
    with host._connection as client:
        _, _, app_alias = test_graceful_shutdown_app.partition("/")

        runners = client.list_runners(start_time=datetime.now() - timedelta(seconds=10))
        runner = next((r for r in runners if r.alias == app_alias), None)

        assert runner is not None, "Runner not found"

        client.kill_runner(runner.runner_id)
    time.sleep(2)

    assert (
        apps.run(test_graceful_shutdown_app, {"uuid": token}, path="/set-uuid") == "ok"
    ), "UUID not set"
    res = apps.run(test_graceful_shutdown_app, {}, path="/latest-request-id")

    teardown_called = res == saved_request_id
    try:
        request_processed = handle.fetch_result()["slept"]
    except Exception:
        request_processed = False

    return teardown_called and request_processed


@pytest.mark.flaky(max_runs=3)
def test_graceful_shutdown(
    host: api.FalServerlessHost,
    test_graceful_shutdown_app: str,
):
    assert graceful_shutdown(
        test_graceful_shutdown_app, host, wait_time=5, path="/"
    ), "app should be gracefully shutdown"


@pytest.mark.flaky(max_runs=3)
def test_graceful_shutdown_force_kill(
    host: api.FalServerlessHost,
    test_graceful_shutdown_app: str,
):
    assert not graceful_shutdown(
        test_graceful_shutdown_app, host, wait_time=60, path="/"
    ), "app should be forcefully killed if it takes too long to clean up"


@pytest.mark.flaky(max_runs=3)
def test_graceful_shutdown_handle_exit(
    host: api.FalServerlessHost,
    test_graceful_shutdown_app: str,
):
    assert graceful_shutdown(
        test_graceful_shutdown_app, host, wait_time=60, path="/with-stop"
    ), "app should be called handle_exit on SIGTERM"


@pytest.mark.flaky(max_runs=3)
def test_runner_machine_type(host: api.FalServerlessHost, test_sleep_app: str):
    """Test that machine_type is populated in runner info."""
    submit_and_wait_for_runner(test_sleep_app, arguments={"wait_time": 1})

    with host._connection as client:
        _, _, app_alias = test_sleep_app.partition("/")

        # list_alias_runners
        runners = client.list_alias_runners(app_alias)
        assert len(runners) >= 1
        assert runners[0].machine_type == "XS"

        # list_runners
        all_runners = client.list_runners(
            start_time=datetime.now() - timedelta(seconds=60)
        )
        assert len(all_runners) >= 1
        target_runner = next((r for r in all_runners if r.alias == app_alias), None)
        assert target_runner is not None, "Runner for test app alias not found"
        assert target_runner.machine_type == "XS"


class RequestContextOutput(BaseModel):
    request_id_from_context: Optional[str]
    endpoint_from_context: Optional[str]
    lifecycle_preference_from_context: Optional[dict]
    request_id_from_header: Optional[str]


def _external_get_request_context() -> dict:
    """External function that accesses request context without request parameter."""
    current_app = get_current_app()
    if current_app is None or current_app.current_request is None:
        return {
            "request_id": None,
            "endpoint": None,
            "lifecycle_preference": None,
        }
    return {
        "request_id": current_app.current_request.request_id,
        "endpoint": current_app.current_request.endpoint,
        "lifecycle_preference": current_app.current_request.lifecycle_preference,
    }


class RequestContextApp(fal.App, keep_alive=300, max_concurrency=1, max_multiplexing=3):
    """App to test request context fields are properly populated."""

    machine_type = "XS"

    @fal.endpoint("/")
    def get_context(self, request: Request) -> RequestContextOutput:
        # Sleep to intentionally cause potential race conditions with multiplexing
        time.sleep(2)

        # Get context from external function (simulates File.from_bytes usage)
        context_data = _external_get_request_context()

        return RequestContextOutput(
            request_id_from_context=context_data["request_id"],
            endpoint_from_context=context_data["endpoint"],
            lifecycle_preference_from_context=context_data["lifecycle_preference"],
            request_id_from_header=request.headers.get("x-fal-request-id"),
        )


@pytest.fixture(scope="module")
def test_request_context_app(host: api.FalServerlessHost, user: User):
    request_context_app = wrap_app(RequestContextApp)
    with register_app(host, request_context_app, "request-context") as (app_alias, _):
        yield f"{user.username}/{app_alias}"


def test_request_context_fields_populated(test_request_context_app: str):
    """Test that request context fields are properly populated."""

    result = apps.run(test_request_context_app, arguments={})

    assert result["request_id_from_context"] is not None
    assert result["endpoint_from_context"] is not None
    assert result["lifecycle_preference_from_context"] is not None

    assert result["request_id_from_context"] == result["request_id_from_header"]


def test_request_context_isolation_with_multiplexing(test_request_context_app: str):
    """Test that request context is properly isolated between concurrent requests.

    With multiplexing enabled (max_multiplexing=3), each request should have
    its own isolated context via ContextVar, ensuring request_id from context
    matches the request_id from headers for each individual request.
    """

    # Submit multiple concurrent requests
    handle_1 = apps.submit(test_request_context_app, arguments={})
    handle_2 = apps.submit(test_request_context_app, arguments={})
    handle_3 = apps.submit(test_request_context_app, arguments={})

    # Get results
    result_1 = handle_1.get()
    result_2 = handle_2.get()
    result_3 = handle_3.get()

    # Each request's context should match its own header, proving isolation
    assert result_1["request_id_from_context"] == result_1["request_id_from_header"]
    assert result_2["request_id_from_context"] == result_2["request_id_from_header"]
    assert result_3["request_id_from_context"] == result_3["request_id_from_header"]

    # All request IDs should be different (unique requests)
    request_ids = {
        result_1["request_id_from_context"],
        result_2["request_id_from_context"],
        result_3["request_id_from_context"],
    }
    assert len(request_ids) == 3, "Each request should have a unique request_id"
