# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# openFuyao is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

from __future__ import annotations

import sys
import tempfile
from pathlib import Path
from unittest import IsolatedAsyncioTestCase

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

from core.errors import NotInitializedError, RenderUnavailableError
from core.models import ChatCompletionRequest, ChatMessage, CompletionRequest, TokenizationResult
from core.service import TokenizerService
from providers.huggingface import HuggingFaceProvider
from providers.resolution import ResolvedTokenizer

REAL_MODEL_ID = "Qwen/Qwen3-8B"
REAL_PROMPT = "hello"


class _FakeRenderClient:
    async def render_completion(self, request: CompletionRequest) -> TokenizationResult:
        del request
        return TokenizationResult(token_ids=[11, 22, 33])

    async def render_chat_completion(self, request) -> TokenizationResult:
        del request
        return TokenizationResult(token_ids=[44, 55, 66])


def _render_factory(resolved: ResolvedTokenizer) -> _FakeRenderClient:
    del resolved
    return _FakeRenderClient()


def _hf_provider(identifier: str) -> HuggingFaceProvider:
    return HuggingFaceProvider(
        ResolvedTokenizer(
            served_model_name=REAL_MODEL_ID,
            model_path=identifier,
            tokenizer_path=identifier,
            provider="huggingface",
            trust_remote_code=True,
        )
    )


class TokenizerServiceTests(IsolatedAsyncioTestCase):
    def test_tokenize_requires_initialize(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )
        with self.assertRaises(NotInitializedError):
            service.tokenize(model=REAL_MODEL_ID, prompt="hello")

    def test_initialize_requires_render_factory(self) -> None:
        service = TokenizerService(provider_name="huggingface")

        with self.assertRaises(RenderUnavailableError):
            service.initialize(model=REAL_MODEL_ID)

    async def test_render_requires_initialize(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )

        with self.assertRaises(NotInitializedError):
            await service.render_completion(
                CompletionRequest(model=REAL_MODEL_ID, prompt_text=REAL_PROMPT)
            )

    async def test_render_completion_rejects_model_mismatch(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )
        service.initialize(model=REAL_MODEL_ID)

        with self.assertRaisesRegex(
            ValueError,
            r"model mismatch: requested='other/model', initialized='Qwen/Qwen3-8B'",
        ):
            await service.render_completion(
                CompletionRequest(model="other/model", prompt_text=REAL_PROMPT)
            )

    async def test_render_chat_completion_rejects_model_mismatch(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )
        service.initialize(model=REAL_MODEL_ID)

        with self.assertRaisesRegex(
            ValueError,
            r"model mismatch: requested='other/model', initialized='Qwen/Qwen3-8B'",
        ):
            await service.render_chat_completion(
                ChatCompletionRequest(
                    model="other/model",
                    messages=[ChatMessage(role="user", content=REAL_PROMPT)],
                )
            )

    def test_tokenize_rejects_model_mismatch(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )
        service.initialize(model=REAL_MODEL_ID)

        with self.assertRaisesRegex(
            ValueError,
            r"model mismatch: requested='other/model', initialized='Qwen/Qwen3-8B'",
        ):
            service.tokenize(model="other/model", prompt=REAL_PROMPT)

    def test_initialize_uses_local_tokenizer_path_as_render_input(self) -> None:
        remote_provider = _hf_provider(REAL_MODEL_ID)
        with tempfile.TemporaryDirectory() as tmp_dir:
            local_dir = Path(tmp_dir) / "local-tokenizer"
            getattr(remote_provider, "_tokenizer").save_pretrained(local_dir)

            service = TokenizerService(
                provider_name="huggingface",
                render_factory=_render_factory,
            )
            result = service.initialize(
                model="my-custom-model",
                tokenizer_name=str(local_dir),
            )

        self.assertEqual(result.model, "my-custom-model")
        self.assertEqual(result.resolved_provider, "huggingface")
        self.assertIsNotNone(service._resolved_tokenizer)
        self.assertEqual(service._resolved_tokenizer.served_model_name, "my-custom-model")
        self.assertEqual(service._resolved_tokenizer.model_path, str(local_dir))
        self.assertEqual(service._resolved_tokenizer.tokenizer_path, str(local_dir))

    def test_initialize_and_tokenize_with_real_provider(self) -> None:
        service = TokenizerService(
            provider_name="huggingface",
            render_factory=_render_factory,
        )
        with tempfile.TemporaryDirectory() as cache_dir:
            service = TokenizerService(
                provider_name="huggingface",
                cache_dir=cache_dir,
                render_factory=_render_factory,
            )
            result = service.initialize(
                model=REAL_MODEL_ID,
            )

        self.assertEqual(result.model, REAL_MODEL_ID)
        self.assertEqual(result.resolved_provider, "huggingface")
        tokenized = service.tokenize(model=REAL_MODEL_ID, prompt=REAL_PROMPT, add_special_tokens=True)
        self.assertTrue(tokenized.token_ids)
        self.assertTrue(all(isinstance(token_id, int) for token_id in tokenized.token_ids))

    def test_render_factory_failure_raises_error(self) -> None:
        def failing_render_factory(resolved: ResolvedTokenizer) -> None:
            del resolved
            raise RuntimeError("render factory failed")

        service = TokenizerService(
            provider_name="huggingface",
            render_factory=failing_render_factory,
        )

        with self.assertRaises(RuntimeError):
            service.initialize(model=REAL_MODEL_ID)

    def test_render_factory_returns_none_raises_error(self) -> None:
        def null_render_factory(resolved: ResolvedTokenizer) -> None:
            del resolved
            return None

        service = TokenizerService(
            provider_name="huggingface",
            render_factory=null_render_factory,
        )

        with self.assertRaises(RenderUnavailableError):
            service.initialize(model=REAL_MODEL_ID)