"""Mistral is a model provider over its own Chat Completions API.

Mistral refuses any request field it does not document and has two shapes of
its own, so the assertions here are on what the daemon sends and how it reads
the answer: no `stream_options` or `max_completion_tokens`, the session as
`prompt_cache_key`, the catalog's reasoning effort, thinking that streams as
typed content chunks and goes back the same way, and tool call ids replayed
as Mistral's nine letters and digits. The picker lists only the models the
API marks as chatting and calling tools, and a key saved by /login serves a
profile that pins none.

The saved key is global, so the class is exclusive.
"""

import json
import unittest
import urllib.parse

from harness import (
    Albedo,
    Provider,
    Reply,
    error,
    exclusive,
    operation_id,
    python,
    text,
)

MODEL = "mistral-medium-latest"

KEY = "mistral-fixture-key"

API_MODELS = {
    "object": "list",
    "data": [
        {
            "id": MODEL,
            "capabilities": {"completion_chat": True, "function_calling": True},
        },
        {
            "id": "mistral-embed",
            "capabilities": {"completion_chat": False, "function_calling": False},
        },
        {
            "id": "mistral-ocr-latest",
            "capabilities": {"completion_chat": True, "function_calling": False},
        },
    ],
}

# models.dev's `mistral` namespace, which names no endpoint: the extension's
# own catalog is what answers for it at a Mistral profile's base url.
MODELS_DEV = {
    "mistral": {
        "id": "mistral",
        "name": "Mistral",
        "env": ["MISTRAL_API_KEY"],
        "models": {
            MODEL: {
                "id": MODEL,
                "limit": {"context": 262144, "output": 262144},
                "reasoning_options": [{"type": "effort", "values": ["none", "high"]}],
            }
        },
    }
}


def chunk(delta, finish=None, usage=None):
    value = {
        "id": "mistral-fixture",
        "object": "chat.completion.chunk",
        "model": MODEL,
        "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
    }
    if usage:
        value["usage"] = usage
    return value


def thinking(text):
    return {"type": "thinking", "thinking": [{"type": "text", "text": text}]}


# A reasoning answer as Mistral streams it: content starts as a list of
# thinking chunks, holds both kinds at the turn, then becomes plain text.
THOUGHTFUL = Reply(
    "text",
    events=[
        # Mistral's delta also carries fields the shared reducer does not read.
        chunk(
            {
                "role": "assistant",
                "content": [thinking("six sevens")],
                "index": 0,
                "tool_call_id": None,
            }
        ),
        chunk({"content": [thinking(" are 42"), {"type": "text", "text": "It is "}]}),
        chunk({"content": "42."}),
        chunk(
            {"content": ""},
            "stop",
            {
                "prompt_tokens": 100,
                "completion_tokens": 10,
                "total_tokens": 110,
                "prompt_tokens_details": {"cached_tokens": 64},
            },
        ),
    ],
)


# A python cell that shows a 1x1 PNG it builds itself.
TINY_PNG = (
    "import struct, zlib\n"
    "def part(kind, data):\n"
    "    body = kind + data\n"
    "    return struct.pack('>I', len(data)) + body + struct.pack('>I', zlib.crc32(body))\n"
    "png = (b'\\x89PNG\\r\\n\\x1a\\n'\n"
    "    + part(b'IHDR', struct.pack('>IIBBBBB', 1, 1, 8, 2, 0, 0, 0))\n"
    "    + part(b'IDAT', zlib.compress(b'\\x00\\xff\\x00\\x00'))\n"
    "    + part(b'IEND', b''))\n"
    "show_image(png)\n"
)


def scripting(replies):
    def reply(request):
        return replies.pop(0) if replies else text("done")

    return reply


@exclusive
class MistralTests(unittest.TestCase):
    def start(self, replies, *, key=True, fixture=None):
        provider = Provider(fixture or scripting(list(replies)), catalog=API_MODELS)
        self.addCleanup(provider.close)
        profile = {
            "extension": "mistral",
            "baseUrl": provider.url,
            "model": MODEL,
            "protocol": "chat_completions",
        }
        if key:
            profile["apiKey"] = KEY

        def offline(app):
            (app.home / "models.json").write_text(json.dumps(MODELS_DEV))

        app = Albedo(provider, providers={"mistral": profile}, prepare=offline)
        app.__enter__()
        self.addCleanup(app.__exit__, None, None, None)
        return app, provider

    def ask(self, app, session, prompt):
        app.prompt(session, prompt).close()
        app.idle(session)

    def test_proxy_drops_empty_assistant_messages_without_breaking_tool_history(self):
        def enforce(request):
            for message in request["messages"]:
                if (
                    message["role"] == "assistant"
                    and not message.get("content")
                    and not message.get("tool_calls")
                ):
                    return error(
                        400, "Assistant message must have content or tool_calls"
                    )
            return text("accepted")

        app, provider = self.start([], fixture=enforce)
        app.write_extensions({"enabled": {"proxy": True}})
        call = {
            "id": "external-call",
            "type": "function",
            "function": {"name": "lookup", "arguments": "{}"},
        }
        history = [
            {"role": "user", "content": "first"},
            {"role": "assistant", "content": None, "refusal": None},
            {"role": "user", "content": "second"},
            {"role": "assistant", "content": "", "tool_calls": []},
            {"role": "assistant", "content": None, "tool_calls": [call]},
            {"role": "tool", "tool_call_id": call["id"], "content": "found"},
            {"role": "assistant", "content": ""},
            {"role": "user", "content": "continue"},
        ]
        with app.api(
            "/extensions/proxy/v1/chat/completions",
            {"model": "mistral/" + MODEL, "messages": history, "stream": True},
        ) as response:
            chunks = [
                json.loads(line.removeprefix("data: "))
                for line in response.read().decode().splitlines()
                if line.startswith("data: ") and line != "data: [DONE]"
            ]
        self.assertFalse([chunk for chunk in chunks if "error" in chunk], chunks)
        self.assertEqual(
            "".join(
                choice["delta"].get("content", "")
                for chunk in chunks
                for choice in chunk.get("choices", [])
            ),
            "accepted",
        )
        [sent] = provider.requests
        messages = sent["request"]["messages"]
        self.assertEqual(
            [message["role"] for message in messages],
            ["user", "user", "assistant", "tool", "assistant", "user"],
        )
        self.assertEqual(messages[0]["content"], "first")
        self.assertEqual(messages[1]["content"], "second")
        self.assertEqual(messages[2]["tool_calls"][0]["function"], call["function"])
        self.assertEqual(
            messages[3]["tool_call_id"], messages[2]["tool_calls"][0]["id"]
        )
        self.assertEqual(messages[4]["content"], "I have processed the tool results.")

    def test_the_request_carries_only_what_mistral_documents(self):
        app, provider = self.start([text("an answer")])
        session = app.session()
        self.ask(app, session, "hello")
        [request] = provider.requests
        self.assertTrue(request["path"].endswith("/chat/completions"), request["path"])
        self.assertEqual(request["authorization"], f"Bearer {KEY}")
        body = request["request"]
        self.assertTrue(body["stream"])
        for absent in ("stream_options", "max_completion_tokens", "store"):
            self.assertNotIn(absent, body)
        self.assertEqual(body["prompt_cache_key"], session)
        # models.dev lists none and high for this model; high is the default.
        self.assertEqual(body["reasoning_effort"], "high")
        self.assertEqual(body["messages"][-1], {"role": "user", "content": "hello"})

    def test_thinking_chunks_stream_and_replay_as_mistral_reads_them(self):
        app, provider = self.start(
            [python("print(6 * 7)", reasoning="let me compute"), THOUGHTFUL]
        )
        session = app.session()
        self.ask(app, session, "multiply")
        answers = [
            entry
            for entry in app.history(session)["items"]
            if entry["kind"] == "assistant"
        ]
        self.assertEqual(
            answers[-1]["content"], [{"kind": "text", "text": "It is 42."}]
        )

        messages = provider.requests[1]["request"]["messages"]
        [call] = [m for m in messages if m["role"] == "assistant"]
        [result] = [m for m in messages if m["role"] == "tool"]
        # The fixture's own id is no Mistral id; its replacement is stable, so
        # the call and its result still pair.
        call_id = call["tool_calls"][0]["id"]
        self.assertRegex(call_id, r"^[A-Za-z0-9]{9}$")
        self.assertEqual(result["tool_call_id"], call_id)
        self.assertIn("42", result["content"])
        self.assertEqual(call["content"], [thinking("let me compute")])
        self.assertEqual(set(call), {"role", "content", "tool_calls"})

        self.ask(app, session, "and doubled?")
        messages = provider.requests[2]["request"]["messages"]
        replayed = [m for m in messages if m["role"] == "assistant"][-1]
        self.assertEqual(
            replayed["content"],
            [thinking("six sevens are 42"), {"type": "text", "text": "It is 42."}],
        )
        for message in messages:
            self.assertFalse(
                {"reasoning_content", "reasoning", "refusal"} & set(message), message
            )

    def test_a_tool_results_images_follow_a_bridge_after_the_results(self):
        # Mistral takes no user message straight after tool results, so the
        # images a result carries come after a short assistant bridge.
        app, provider = self.start([python(TINY_PNG), text("a dot")])
        self.ask(app, app.session(), "draw")
        messages = provider.requests[1]["request"]["messages"]
        roles = [m["role"] for m in messages]
        tool = roles.index("tool")
        self.assertEqual(roles[tool + 1 :], ["assistant", "user"])
        self.assertEqual(
            messages[tool + 1]["content"], "I have processed the tool results."
        )
        self.assertIsInstance(messages[tool]["content"], str)
        parts = messages[tool + 2]["content"]
        self.assertEqual(
            parts[0]["text"],
            f"Images from tool call {messages[tool]['tool_call_id']}:",
        )
        self.assertTrue(parts[1]["image_url"]["url"].startswith("data:image/png"))

    def test_the_picker_lists_the_chat_models_that_call_tools(self):
        app, provider = self.start([])
        query = urllib.parse.urlencode(
            {"provider": "mistral", "endpoint": provider.url}
        )
        with app.api(f"/models?{query}") as response:
            listed = [item["id"] for item in json.load(response)["items"]]
        self.assertEqual(listed, [MODEL])

    def test_a_key_saved_by_login_serves_a_profile_that_pins_none(self):
        app, provider = self.start([text("keyed")], key=False)
        with app.api(
            f"/auth/keys/{operation_id()}",
            {
                "provider": "mistral",
                "endpoint": provider.url,
                "api_key": KEY,
                "label": "fixture",
            },
            method="PUT",
        ) as response:
            self.assertEqual(json.load(response)["kind"], "key")
        self.ask(app, app.session(), "hello")
        self.assertEqual(provider.requests[0]["authorization"], f"Bearer {KEY}")


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