from __future__ import annotations

import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import MagicMock, patch


BACKEND_DIR = Path(__file__).resolve().parents[1]
os.environ["KEFU_DB_TYPE"] = "sqlite"
os.environ["KEFU_DATABASE_URL"] = ""
os.environ["KEFU_DB_PATH"] = str(BACKEND_DIR / "kefu.db")
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

import main


class _FakeResponse:
    def __init__(
        self,
        *,
        final_host: str,
        headers: dict[str, str] | None = None,
        chunks: tuple[bytes, ...] = (),
    ) -> None:
        self.url = SimpleNamespace(host=final_host)
        self.headers = headers or {}
        self._chunks = chunks
        self.raise_for_status = MagicMock()
        self.iteration_started = False

    async def aiter_bytes(self):
        self.iteration_started = True
        for chunk in self._chunks:
            yield chunk


class _AsyncContext:
    def __init__(self, value) -> None:
        self.value = value

    async def __aenter__(self):
        return self.value

    async def __aexit__(self, exc_type, exc, traceback):
        return False


class _FakeAsyncClient:
    def __init__(self, response: _FakeResponse) -> None:
        self.response = response
        self.stream_calls: list[tuple[tuple, dict]] = []

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, traceback):
        return False

    def stream(self, *args, **kwargs):
        self.stream_calls.append((args, kwargs))
        return _AsyncContext(self.response)


class MediaProxyHostGuardTests(unittest.TestCase):
    def test_allows_each_root_domain_and_its_subdomains(self):
        for allowed in main._MEDIA_PROXY_HOSTS:
            with self.subTest(host=allowed):
                self.assertTrue(main._is_allowed_media_host(allowed))
                self.assertTrue(main._is_allowed_media_host(f"cdn.images.{allowed}"))

        self.assertTrue(main._is_allowed_media_host("CDN.DOUYINPIC.COM."))

    def test_rejects_empty_suffix_tricks_and_similar_domains(self):
        rejected = (
            "",
            ".",
            "douyin.com.evil",
            "evildouyin.com",
            "byteimg.com.evil.example",
            "evilbyteimg.com",
            "ibyteimg.comevil",
            "douyin.co",
            "douyinpic.co",
            "douyin-static.com",
            "snssdk.example",
        )

        for hostname in rejected:
            with self.subTest(host=hostname):
                self.assertFalse(main._is_allowed_media_host(hostname))


class MediaProxyResponseGuardTests(unittest.IsolatedAsyncioTestCase):
    async def _call_proxy(self, response: _FakeResponse):
        client = _FakeAsyncClient(response)
        constructor = MagicMock(return_value=client)
        with patch.object(main.httpx, "AsyncClient", constructor):
            result = await main.proxy_media(
                url="https://cdn.douyinpic.com/media/test.jpg"
            )
        return result, client, constructor

    async def test_rejects_redirect_to_non_allowlisted_final_domain(self):
        response = _FakeResponse(
            final_host="attacker.example",
            headers={"content-type": "image/jpeg"},
            chunks=(b"not-read",),
        )

        with self.assertRaises(main.HTTPException) as caught:
            await self._call_proxy(response)

        self.assertEqual(caught.exception.status_code, 400)
        self.assertFalse(response.iteration_started)

    async def test_rejects_oversized_content_length_before_streaming(self):
        response = _FakeResponse(
            final_host="cdn.douyinpic.com",
            headers={"content-length": "9", "content-type": "image/jpeg"},
            chunks=(b"not-read",),
        )

        with (
            patch.object(main, "_MEDIA_PROXY_MAX_BYTES", 8),
            self.assertRaises(main.HTTPException) as caught,
        ):
            await self._call_proxy(response)

        self.assertEqual(caught.exception.status_code, 413)
        self.assertFalse(response.iteration_started)

    async def test_rejects_stream_when_actual_bytes_exceed_limit(self):
        response = _FakeResponse(
            final_host="cdn.douyinpic.com",
            headers={"content-length": "0", "content-type": "image/jpeg"},
            chunks=(b"12345", b"6789"),
        )

        with (
            patch.object(main, "_MEDIA_PROXY_MAX_BYTES", 8),
            self.assertRaises(main.HTTPException) as caught,
        ):
            await self._call_proxy(response)

        self.assertEqual(caught.exception.status_code, 413)
        self.assertTrue(response.iteration_started)


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