from __future__ import annotations

import asyncio
import base64
import json
import unittest
from unittest.mock import AsyncMock, patch

from fastapi import HTTPException
from routes import voice_webrtc
from starlette.requests import Request


class FakeConnection:
    pc_id = "pc_test"


def encoded_header_json(value: object) -> str:
    payload = json.dumps(value).encode("utf-8")
    return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=")


def offer_request(*, client_tools: object | None = None) -> Request:
    headers = []
    if client_tools is not None:
        headers.append(
            (
                b"x-pipecat-client-tools",
                encoded_header_json(client_tools).encode("ascii"),
            )
        )
    return Request({"type": "http", "method": "POST", "headers": headers})


class WebRTCOfferMetadataTest(unittest.IsolatedAsyncioTestCase):
    async def test_client_tool_header_survives_request_endpoint_body_replacement(self):
        request = offer_request(
            client_tools=["set_photo_button_visible", "show_message"]
        )
        payload = {"sdp": "offer", "type": "offer"}

        with patch.object(
            voice_webrtc,
            "_handle_offer_payload",
            new=AsyncMock(return_value={"type": "answer"}),
        ) as handle_offer:
            await voice_webrtc.http_offer(request, payload)

        forwarded = handle_offer.await_args.args[0]
        self.assertEqual(
            forwarded["client_tool_names"],
            ["set_photo_button_visible", "show_message"],
        )

    async def test_client_tool_request_data_remains_a_compatibility_fallback(self):
        request = offer_request()
        payload = {
            "sdp": "offer",
            "type": "offer",
            "requestData": {"client_tool_names": ["show_message"]},
        }

        with patch.object(
            voice_webrtc,
            "_handle_offer_payload",
            new=AsyncMock(return_value={"type": "answer"}),
        ) as handle_offer:
            await voice_webrtc.http_offer(request, payload)

        forwarded = handle_offer.await_args.args[0]
        self.assertEqual(forwarded["client_tool_names"], ["show_message"])

    async def test_invalid_client_tool_header_is_rejected(self):
        request = offer_request(client_tools={"name": "show_message"})

        with self.assertRaises(HTTPException) as raised:
            await voice_webrtc.http_offer(
                request,
                {"sdp": "offer", "type": "offer"},
            )

        self.assertEqual(raised.exception.status_code, 400)


class WebRTCPipelineLifecycleTest(unittest.IsolatedAsyncioTestCase):
    async def asyncTearDown(self):
        tasks = list(voice_webrtc._pipeline_tasks)
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        voice_webrtc._pipeline_tasks.clear()
        voice_webrtc._connection_tasks.clear()
        voice_webrtc._active_connections.clear()

    async def test_pipeline_task_is_retained_until_completion(self):
        connection = FakeConnection()
        release = asyncio.Event()

        async def pipeline():
            await release.wait()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        self.assertIn(task, voice_webrtc._pipeline_tasks)
        self.assertIs(voice_webrtc._connection_tasks[connection], task)

        release.set()
        await task
        await asyncio.sleep(0)

        self.assertNotIn(task, voice_webrtc._pipeline_tasks)
        self.assertNotIn(connection, voice_webrtc._connection_tasks)

    async def test_stuck_pipeline_is_cancelled_after_close_grace_period(self):
        connection = FakeConnection()
        cancelled = asyncio.Event()

        async def pipeline():
            try:
                await asyncio.Event().wait()
            finally:
                cancelled.set()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        with patch.object(voice_webrtc, "PIPELINE_CLOSE_GRACE_SECONDS", 0.01):
            await voice_webrtc._wait_for_pipeline_close(
                task,
                connection_id=connection.pc_id,
            )

        self.assertTrue(cancelled.is_set())
        self.assertTrue(task.cancelled())


if __name__ == "__main__":
    unittest.main()