# Copyright 2022-2026 Xinference Holdings Pte. Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import json
import os
import time
from concurrent.futures import ThreadPoolExecutor

import psutil
import pytest
import requests

from ...constants import XINFERENCE_ENV_MODEL_SRC
from ..restful.restful_client import Client as RESTfulClient
from ..restful.restful_client import (
    RESTfulChatModelHandle,
    RESTfulEmbeddingModelHandle,
    _get_error_string,
)


class _DummyResponse:
    status_code = 200

    def json(self):
        return {"choices": [{"message": {"content": "ok"}}]}


class _DummySession:
    def __init__(self):
        self.last_json = None

    def post(self, url, json=None, stream=None, headers=None):
        self.last_json = json
        return _DummyResponse()

    def close(self):
        return None


@pytest.mark.skipif(os.name == "nt", reason="Skip windows")
def test_RESTful_client(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)
    assert len(client.list_models()) == 0

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    assert len(client.list_models()) == 1

    model = client.get_model(model_uid=model_uid)
    assert isinstance(model, RESTfulChatModelHandle)

    with pytest.raises(RuntimeError):
        model = client.get_model(model_uid="test")
        assert isinstance(model, RESTfulChatModelHandle)

    with pytest.raises(RuntimeError):
        completion = model.generate({"max_tokens": 64})

    completion = model.generate("Once upon a time, there was a very old computer")
    assert "text" in completion["choices"][0]

    completion = model.generate(
        "Once upon a time, there was a very old computer", {"max_tokens": 64}
    )
    assert "text" in completion["choices"][0]

    streaming_response = model.generate(
        "Once upon a time, there was a very old computer",
        {"max_tokens": 64, "stream": True},
    )

    for chunk in streaming_response:
        assert "text" in chunk["choices"][0]

    with pytest.raises(RuntimeError):
        completion = model.chat({"max_tokens": 64})

    messages = [{"role": "user", "content": "What is the capital of France?"}]
    completion = model.chat(messages)
    assert "content" in completion["choices"][0]["message"]

    def _check_stream():
        streaming_response = model.chat(
            messages,
            generate_config={"stream": True, "max_tokens": 5},
        )
        for chunk in streaming_response:
            if not chunk["choices"]:
                assert chunk["usage"]
                continue
            assert "finish_reason" in chunk["choices"][0]
            finish_reason = chunk["choices"][0]["finish_reason"]
            if finish_reason is None:
                assert ("content" in chunk["choices"][0]["delta"]) or (
                    "role" in chunk["choices"][0]["delta"]
                )
            else:
                assert chunk["choices"][0]["delta"] == {"content": ""}

    _check_stream()

    results = []
    with ThreadPoolExecutor() as executor:
        for _ in range(2):
            r = executor.submit(_check_stream)
            results.append(r)

    for r in results:
        r.result()

    # After iteration finish, we can iterate again.
    _check_stream()

    client.terminate_model(model_uid=model_uid)
    assert len(client.list_models()) == 0

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    assert len(client.list_models()) == 1

    # Test concurrent chat is OK.
    model = client.get_model(model_uid=model_uid)

    def _check(stream=False):
        completion = model.generate(
            "AI is going to", generate_config={"stream": stream, "max_tokens": 5}
        )
        if stream:
            count = 0
            has_text = False
            for chunk in completion:
                assert "text" in chunk["choices"][0]
                if chunk["choices"][0]["text"]:
                    has_text = True
                count += 1
            assert has_text
            assert count > 2
        else:
            assert "text" in completion["choices"][0]
            assert len(completion["choices"][0]["text"]) > 0

    for stream in [True, False]:
        results = []
        with ThreadPoolExecutor() as executor:
            for _ in range(3):
                r = executor.submit(_check, stream=stream)
                results.append(r)
        for r in results:
            r.result()

    client.terminate_model(model_uid=model_uid)
    assert len(client.list_models()) == 0

    with pytest.raises(RuntimeError):
        client.terminate_model(model_uid=model_uid)

    model_uid2 = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )

    client.terminate_model(model_uid=model_uid2)
    assert len(client.list_models()) == 0


def test_query_engines_by_name(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)

    assert len(client.query_engine_by_model_name("qwen3")) > 0
    assert len(client.query_engine_by_model_name("qwen3", model_type=None)) > 0
    assert len(client.query_engine_by_model_name("bge-m3", model_type="embedding")) > 0


@pytest.mark.skipif(os.name == "nt", reason="Skip windows")
def test_list_cached_models(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)
    res = client.list_cached_models()
    assert len(res) > 0


def test_RESTful_client_for_embedding(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)
    assert len(client.list_models()) == 0

    model_uid = client.launch_model(model_name="gte-base", model_type="embedding")
    assert len(client.list_models()) == 1

    model = client.get_model(model_uid=model_uid)
    assert isinstance(model, RESTfulEmbeddingModelHandle)

    completion = model.create_embedding("write a poem.")
    assert len(completion["data"][0]["embedding"]) == 768

    client.terminate_model(model_uid=model_uid)
    assert len(client.list_models()) == 0


def test_RESTful_client_custom_model(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)

    model_regs = client.list_model_registrations(model_type="LLM")
    assert len(model_regs) > 0
    for model_reg in model_regs:
        assert model_reg["is_builtin"]

    model = """{
  "version": 2,
  "context_length":2048,
  "model_name": "custom_model",
  "model_lang": [
    "en", "zh"
  ],
  "model_ability": [
    "chat"
  ],
  "model_family": "other",
  "model_specs": [
    {
      "model_format": "pytorch",
      "model_size_in_billions": 7,
      "quantization": "none",
      "model_id": "ziqingyang/chinese-alpaca-2-7b"
    }
  ],
  "chat_template": "xyz",
  "stop_token_ids": [],
  "stop": []
}"""
    client.register_model(model_type="LLM", model=model, persist=False)

    data = client.get_model_registration(model_type="LLM", model_name="custom_model")
    assert "custom_model" in data["model_name"]

    new_model_regs = client.list_model_registrations(model_type="LLM")
    assert len(new_model_regs) == len(model_regs) + 1
    custom_model_reg = None
    for model_reg in new_model_regs:
        if model_reg["model_name"] == "custom_model":
            custom_model_reg = model_reg
    assert custom_model_reg is not None

    client.unregister_model(model_type="LLM", model_name="custom_model")
    new_model_regs = client.list_model_registrations(model_type="LLM")
    assert len(new_model_regs) == len(model_regs)
    custom_model_reg = None
    for model_reg in new_model_regs:
        if model_reg["model_name"] == "custom_model":
            custom_model_reg = model_reg
    assert custom_model_reg is None

    # test register with chat_template using model_family
    model_with_prompt = """{
  "version": 2,
  "context_length":2048,
  "model_name": "custom_model",
  "model_lang": [
    "en", "zh"
  ],
  "model_ability": [
    "embed",
    "chat"
  ],
  "model_family": "other",
  "model_specs": [
    {
      "model_format": "pytorch",
      "model_size_in_billions": 7,
      "quantization": "none",
      "model_id": "ziqingyang/chinese-alpaca-2-7b"
    }
  ],
  "chat_template": "qwen-chat"
}"""
    client.register_model(model_type="LLM", model=model_with_prompt, persist=False)
    client.unregister_model(model_type="LLM", model_name="custom_model")

    model_with_vision = """{
      "version": 2,
      "context_length":2048,
      "model_name": "custom_model",
      "model_lang": [
        "en", "zh"
      ],
      "model_ability": [
        "chat",
        "vision"
      ],
      "model_family": "other",
      "model_specs": [
        {
          "model_format": "pytorch",
          "model_size_in_billions": 7,
          "quantization": "none",
          "model_id": "ziqingyang/chinese-alpaca-2-7b"
        }
      ],
      "chat_template": "xyz123"
    }"""
    with pytest.raises(RuntimeError):
        client.register_model(model_type="LLM", model=model_with_vision, persist=False)

    model_with_tool_call = """{
          "version": 2,
          "context_length":2048,
          "model_name": "custom_model",
          "model_lang": [
            "en", "zh"
          ],
          "model_ability": [
            "chat",
            "tools"
          ],
          "model_family": "other",
          "model_specs": [
            {
              "model_format": "pytorch",
              "model_size_in_billions": 7,
              "quantization": "none",
              "model_id": "ziqingyang/chinese-alpaca-2-7b"
            }
          ],
          "chat_template": "xyz123"
        }"""
    with pytest.raises(RuntimeError):
        client.register_model(
            model_type="LLM", model=model_with_tool_call, persist=False
        )


def test_client_from_modelscope(setup):
    try:
        os.environ[XINFERENCE_ENV_MODEL_SRC] = "modelscope"

        endpoint, _ = setup
        client = RESTfulClient(endpoint)
        assert len(client.list_models()) == 0

        model_uid = client.launch_model(
            model_name="tiny-llama", model_engine="llama.cpp"
        )
        assert len(client.list_models()) == 1
        model = client.get_model(model_uid=model_uid)
        completion = model.generate("write a poem.", generate_config={"max_tokens": 5})
        assert "text" in completion["choices"][0]
        assert len(completion["choices"][0]["text"]) > 0
    finally:
        os.environ.pop(XINFERENCE_ENV_MODEL_SRC)


def test_client_error():
    r = requests.Response()
    r.url = "0.0.0.0:1234"
    r.status_code = 502
    r.reason = "Bad Gateway"
    err = _get_error_string(r)
    assert "502 Server Error: Bad Gateway for url: 0.0.0.0:1234" == err
    r._content = json.dumps({"detail": "Test error"}).encode("utf-8")
    err = _get_error_string(r)
    assert "Test error" == err


def test_client_custom_embedding_model(setup):
    endpoint, _ = setup
    client = RESTfulClient(endpoint)

    model_regs = client.list_model_registrations(model_type="embedding")
    assert len(model_regs) > 0
    for model_reg in model_regs:
        assert model_reg["is_builtin"]

    model = """{
  "model_name": "custom-bge-small-en",
  "dimensions": 1024,
  "max_tokens": 512,
  "language": ["en"],
  "model_specs": [
    {
      "model_format": "pytorch",
      "model_id": "Xorbits/bge-small-en",
      "quantization": "none"
    }
  ]
}"""
    client.register_model(model_type="embedding", model=model, persist=False)

    data = client.get_model_registration(
        model_type="embedding", model_name="custom-bge-small-en"
    )
    assert "custom-bge-small-en" in data["model_name"]

    new_model_regs = client.list_model_registrations(model_type="embedding")
    assert len(new_model_regs) == len(model_regs) + 1
    custom_model_reg = None
    for model_reg in new_model_regs:
        if model_reg["model_name"] == "custom-bge-small-en":
            custom_model_reg = model_reg
    assert custom_model_reg is not None
    assert not custom_model_reg["is_builtin"]

    # unregister
    client.unregister_model(model_type="embedding", model_name="custom-bge-small-en")
    new_model_regs = client.list_model_registrations(model_type="embedding")
    assert len(new_model_regs) == len(model_regs)
    custom_model_reg = None
    for model_reg in new_model_regs:
        if model_reg["model_name"] == "custom-bge-small-en":
            custom_model_reg = model_reg
    assert custom_model_reg is None


@pytest.fixture
def set_auto_recover_limit():
    os.environ["XINFERENCE_MODEL_ACTOR_AUTO_RECOVER_LIMIT"] = "1"
    yield
    del os.environ["XINFERENCE_MODEL_ACTOR_AUTO_RECOVER_LIMIT"]


@pytest.fixture
def set_test_oom_error():
    os.environ["XINFERENCE_TEST_OUT_OF_MEMORY_ERROR"] = "1"
    yield
    del os.environ["XINFERENCE_TEST_OUT_OF_MEMORY_ERROR"]


@pytest.fixture
def setup_cluster():
    import xoscar as xo

    from ...api.restful_api import run_in_subprocess as restful_api_run_in_subprocess
    from ...conftest import TEST_FILE_LOGGING_CONF, TEST_LOGGING_CONF, api_health_check
    from ...deploy.local import health_check
    from ...deploy.local import run_in_subprocess as supervisor_run_in_subprocess

    # This fixture is used by tests that exercise unauthenticated requests;
    # advanced auth defaults to on, so it must be explicitly disabled here,
    # before any subprocess (which inherits this env) is started.
    os.environ["XINFERENCE_AUTH_ADVANCED"] = "false"

    supervisor_address = f"localhost:{xo.utils.get_next_port()}"
    local_cluster = supervisor_run_in_subprocess(
        supervisor_address, None, None, TEST_LOGGING_CONF
    )

    if not health_check(address=supervisor_address, max_attempts=20, sleep_interval=1):
        raise RuntimeError("Supervisor is not available after multiple attempts")

    try:
        port = xo.utils.get_next_port()
        restful_api_proc = restful_api_run_in_subprocess(
            supervisor_address,
            host="localhost",
            port=port,
            logging_conf=TEST_FILE_LOGGING_CONF,
        )
        endpoint = f"http://localhost:{port}"
        if not api_health_check(endpoint, max_attempts=10, sleep_interval=5):
            raise RuntimeError("Endpoint is not available after multiple attempts")

        yield f"http://localhost:{port}", supervisor_address
        restful_api_proc.kill()
    finally:
        local_cluster.kill()


def test_auto_recover(set_auto_recover_limit, setup_cluster):
    endpoint, _ = setup_cluster
    current_proc = psutil.Process()
    chilren_proc = set(current_proc.children(recursive=True))
    client = RESTfulClient(endpoint)

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    new_children_proc = set(current_proc.children(recursive=True))
    model_proc = next(iter(new_children_proc - chilren_proc))
    assert len(client.list_models()) == 1

    model = client.get_model(model_uid=model_uid)
    assert isinstance(model, RESTfulChatModelHandle)

    completion = model.generate("Once upon a time, there was a very old computer")
    assert "text" in completion["choices"][0]

    model_proc.kill()

    for _ in range(60):
        try:
            completion = model.generate(
                "Once upon a time, there was a very old computer", {"max_tokens": 64}
            )
            assert "text" in completion["choices"][0]
            break
        except Exception:
            time.sleep(1)
    else:
        assert False


def test_model_error(set_test_oom_error, setup_cluster):
    endpoint, _ = setup_cluster
    client = RESTfulClient(endpoint)

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    assert len(client.list_models()) == 1

    model = client.get_model(model_uid=model_uid)
    assert isinstance(model, RESTfulChatModelHandle)

    with pytest.raises(RuntimeError, match="Model actor is out of memory"):
        model.generate("Once upon a time, there was a very old computer")


def test_model_oom_triggers_restart(set_test_oom_error, setup_cluster):
    # OOM must hard-exit the model subprocess via os._exit(1) so xoscar
    # can detect the exit and recover it. With the old sys.exit(1) the
    # SystemExit was swallowed by xoscar's _ErrorProcessor, the process
    # never died, and this test would time out at the "subprocess did not
    # exit" assert below.
    endpoint, _ = setup_cluster
    current_proc = psutil.Process()
    before = set(current_proc.children(recursive=True))
    client = RESTfulClient(endpoint)

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    assert len(client.list_models()) == 1
    model_proc = next(iter(set(current_proc.children(recursive=True)) - before))

    model = client.get_model(model_uid=model_uid)
    assert isinstance(model, RESTfulChatModelHandle)

    # Trigger OOM. The handler runs os._exit(1) on the model subprocess, so the
    # client sees an error (OOM report or connection drop) -- either is fine.
    with pytest.raises(Exception):
        model.generate("Once upon a time, there was a very old computer")

    # The old subprocess must actually exit (this is the regression guard).
    for _ in range(30):
        if not model_proc.is_running() or model_proc.status() == psutil.STATUS_ZOMBIE:
            break
        time.sleep(1)
    else:
        assert False, "OOM subprocess did not exit (sys.exit swallowed by xoscar)"


def test_model_oom_recover_exhausted_evicts_replica(
    set_auto_recover_limit, set_test_oom_error, setup_cluster
):
    # Once worker exhausts AUTO_RECOVER_LIMIT (=1), it stops recreating and
    # calls back supervisor.mark_replica_dead, which evicts the dead replica
    # from the round-robin scheduler. For a single-replica model that means
    # the model is taken fully offline, so list_models() drains to empty.
    endpoint, _ = setup_cluster
    client = RESTfulClient(endpoint)

    model_uid = client.launch_model(
        model_name="qwen2.5-instruct",
        model_engine="llama.cpp",
        model_size_in_billions="0_5",
        model_format="ggufv2",
        quantization="q4_0",
    )
    assert len(client.list_models()) == 1

    evicted = False
    for _ in range(60):
        try:
            model = client.get_model(model_uid=model_uid)
            model.generate("Once upon a time, there was a very old computer")
        except Exception:
            pass
        if len(client.list_models()) == 0:
            evicted = True
            break
        time.sleep(1)
    assert (
        evicted
    ), "dead replica was not evicted from supervisor after recover exhausted"


def test_restful_chat_enable_thinking_injected():
    handle = RESTfulChatModelHandle("test-model", "http://localhost", {})
    dummy_session = _DummySession()
    handle.session = dummy_session

    messages = [{"role": "user", "content": "hi"}]
    handle.chat(messages, enable_thinking=False)

    assert dummy_session.last_json["chat_template_kwargs"] == {
        "enable_thinking": False,
        "thinking": False,
    }
    handle.session = None
