import asyncio
import json
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, Generator, Iterator, List, 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: int | str | float


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(0.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()


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"]


# 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"
    teardown_file = "/data/teardown.txt"
    latest_request_id = None

    def setup(self):
        self.stop = False

    def handle_exit(self):
        self.stop = True

    @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:
        with open(self.teardown_file) as f:
            return f.read()

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

        with open(self.teardown_file, "w") as f:
            print("teardown called, latest request id: ", self.latest_request_id)
            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)

    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)

    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"


class RequestContextOutput(BaseModel):
    request_id_from_context: str | None
    endpoint_from_context: str | None
    lifecycle_preference_from_context: dict | None
    request_id_from_header: str | None


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"
