"""DeepSeek API のローカルモックサーバー。

DeepSeek の公式ドキュメント（Search & Fetch / Thinking Mode / Tool Calls の各ガイド）が
定める挙動を、標準ライブラリだけで再現する。API キーが無くても 01〜05 を動かせる。

起動:
    python mock_server.py --port 8787

テストからの利用:
    server, port = make_server(port=0)
    thread = serve_in_thread(server)
    ...
    server.shutdown()
    server.server_close()
"""

from __future__ import annotations

import argparse
import json
import re
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any

KNOWN_MODELS = {"deepseek-flash", "deepseek-v4-pro"}


def _usage_from_text(prompt_text: str, completion_text: str) -> dict[str, int]:
    # トークン数は概算（文字数 / 2 を四捨五入、最低 1）。モックなので厳密さは不要。
    prompt_tokens = max(1, round(len(prompt_text) / 2))
    completion_tokens = max(1, round(len(completion_text) / 2))
    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
        "prompt_cache_hit_tokens": 0,
        "prompt_cache_miss_tokens": prompt_tokens,
    }


def _make_search_result(index: int, query: str, need_content: bool) -> dict[str, Any]:
    # 公式ガイド（Search & Fetch）の例に合わせ、rerank_score / authority は固定の降順系列、1 件だけ date が null。
    rerank_scores = [2.6, 2.1, 1.7, 1.2, 0.8, 0.4]
    authorities = [4, 3, 3, 2, -1, 1]
    score = rerank_scores[index]
    authority = authorities[index]
    date = None if index == 4 else f"2026-10-{10 + index:02d}T00:00:00+00:00"
    summary = (
        f"これは「{query}」に関するモックの要約です。"
        "検索結果の本文を要約した体で、モデルに読ませる用途を想定しています。"
        * 6
    )[:620]
    return {
        "rerank_score": score,
        "title": f"モック結果 {index + 1}: {query}",
        "url": f"https://example.com/mock/{index + 1}",
        "snippet": f"「{query}」に関する短い要約（モック、{index + 1} 件目）。",
        "summary": summary,
        "content": (f"# モック結果 {index + 1}\n\n「{query}」の本文（モック）。" if need_content else ""),
        "site_name": f"Example Site {index + 1}",
        "site_icon": f"https://example.com/favicon{index + 1}.ico",
        "date": date,
        "authority": authority,
    }


_URL_RE = re.compile(r"^(https?)://([^/@\s]+@)?([^/\s]+)(/.*)?$")


def _validate_fetch_url(url: Any) -> str | None:
    """問題があればエラーメッセージを返す。問題なければ None。"""
    if not isinstance(url, str) or not url:
        return "'url' is required"
    if len(url) > 8192:
        return "'url' must be at most 8192 characters"
    m = _URL_RE.match(url)
    if not m:
        return "'url' must be an http or https URL with a host"
    if m.group(2):
        return "'url' must not contain a username or password"
    return None


class MockHandler(BaseHTTPRequestHandler):
    server: "MockServer"

    def log_message(self, format: str, *args: Any) -> None:  # noqa: A002
        pass  # テスト実行中のコンソールを汚さない

    def _send_json(self, status: int, payload: dict[str, Any]) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _record(self, path: str, body: Any) -> None:
        self.server.requests.append(
            {
                "method": "POST",
                "path": path,
                "headers": {k.lower(): v for k, v in self.headers.items()},
                "body": body,
            }
        )

    def do_POST(self) -> None:  # noqa: N802
        try:
            self._dispatch()
        except Exception as exc:  # 想定外の例外でソケットを落とさず 500 を返す
            self._send_json(500, {"error": {"message": str(exc), "type": "internal_error", "param": None, "code": "internal_error"}})

    def _dispatch(self) -> None:
        length = int(self.headers.get("Content-Length", "0") or "0")
        raw = self.rfile.read(length) if length else b""

        auth = self.headers.get("Authorization", "")
        if not auth.startswith("Bearer ") or not auth[len("Bearer "):].strip():
            self._send_json(401, {"error": {"message": "missing or invalid API key", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
            return

        try:
            body = json.loads(raw) if raw else {}
        except json.JSONDecodeError:
            self._send_json(400, {"error": {"message": "invalid JSON body", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
            return

        self._record(self.path, body)

        if self.path == "/v1/tools/search":
            self._handle_search(body)
        elif self.path == "/v1/tools/fetch":
            self._handle_fetch(body)
        elif self.path == "/chat/completions":
            self._handle_chat(body)
        else:
            self._send_json(404, {"error": {"message": "not found", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})

    def _handle_search(self, body: dict[str, Any]) -> None:
        query = body.get("query")
        if not isinstance(query, str) or not (1 <= len(query) <= 128) or not query.strip():
            self._send_json(400, _error_body_dict("'query' must be between 1 and 128 characters"))
            return
        count = body.get("count", 10)
        if isinstance(count, bool) or not isinstance(count, int) or not (1 <= count <= 50):
            self._send_json(400, _error_body_dict("'count' must be between 1 and 50"))
            return
        ori_question = body.get("ori_question")
        if ori_question is not None and (not isinstance(ori_question, str) or not (1 <= len(ori_question) <= 128)):
            self._send_json(400, _error_body_dict("'ori_question' must be between 1 and 128 characters"))
            return
        need_content = body.get("need_content", False)
        if not isinstance(need_content, bool):
            self._send_json(400, _error_body_dict("'need_content' must be a boolean"))
            return

        if "mock:500" in query:
            self._send_json(500, {"error": {"message": "internal error (mock)", "type": "internal_error", "param": None, "code": "internal_error"}})
            return
        if "mock:402" in query:
            self._send_json(402, {"error": {"message": "Insufficient Balance", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
            return
        if "mock:429" in query:
            counts = self.server.search_429_counts
            seen = counts.get(query, 0)
            counts[query] = seen + 1
            if seen < 2:
                self._send_json(429, {"error": {"message": "QPS limit exceeded (mock)", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
                return
            # 3 回目は通常どおり成功させる。
        if "mock:408" in query:
            counts = self.server.search_408_counts
            seen = counts.get(query, 0)
            counts[query] = seen + 1
            if seen < 2:
                self._send_json(408, {"error": {"message": "request body read timed out (mock)", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
                return
            # 3 回目は通常どおり成功させる。
        if "mock:empty" in query:
            self._send_json(200, {"results": []})
            return

        n = min(count, 6)
        results = [_make_search_result(i, query, need_content) for i in range(n)]
        self._send_json(200, {"results": results})

    def _handle_fetch(self, body: dict[str, Any]) -> None:
        url = body.get("url")
        error = _validate_fetch_url(url)
        if error:
            self._send_json(400, _error_body_dict(error))
            return
        m = _URL_RE.match(url)
        host = m.group(3)
        if host == "blocked.example":
            self._send_json(403, {"error": {"message": "rejected by security policy", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
            return
        if host == "missing.example":
            self._send_json(404, {"error": {"message": "could not retrieve usable page content", "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}})
            return
        self._send_json(200, {"title": "Mock page", "content": f"# Mock page\n\nこれは {url} のモック本文です。"})

    def _handle_chat(self, body: dict[str, Any]) -> None:
        model = body.get("model")
        if model not in KNOWN_MODELS:
            self._send_json(400, _error_body_dict(f"unknown model '{model}'"))
            return

        thinking = body.get("thinking") or {}
        is_thinking = thinking.get("type") != "disabled"
        messages = body.get("messages") or []
        tools = body.get("tools")

        # 演習用: 質問文に mock:chat400 が含まれていれば、クライアント側のエラー処理
        # （openai の例外をどう表示するか）を確かめるためにわざと 400 を返す。
        if any("mock:chat400" in str(message.get("content", "")) for message in messages):
            self._send_json(400, _error_body_dict("forced 400 for testing (mock:chat400)"))
            return

        if tools and is_thinking:
            # 公式ガイド（Thinking Mode）: tools 付きの要求では、過去の全 assistant ターンの
            # reasoning_content を返送する必要がある。公式テキストは「返送する必要がある」としか
            # 書いておらず、空文字を拒否するとは書いていないため、このモックはキーの有無だけを見る
            # （これはモックによる解釈であり、公式ドキュメントがここまで明記しているわけではない）。
            for message in messages:
                if message.get("role") == "assistant" and "reasoning_content" not in message:
                    self._send_json(
                        400,
                        _error_body_dict(
                            "reasoning_content of previous assistant messages must be passed back when tools are used"
                        ),
                    )
                    return

        last = messages[-1] if messages else {}
        content, reasoning_content, tool_calls, finish_reason = self._build_chat_reply(
            messages, last, tools, is_thinking
        )

        message: dict[str, Any] = {"role": "assistant", "content": content}
        if is_thinking:
            message["reasoning_content"] = reasoning_content
        if tool_calls:
            message["tool_calls"] = tool_calls

        prompt_text = json.dumps(messages, ensure_ascii=False)
        usage = _usage_from_text(prompt_text, content + (reasoning_content or ""))

        self._send_json(
            200,
            {
                "id": "mock-chatcmpl-1",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model,
                "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
                "usage": usage,
            },
        )

    def _build_chat_reply(
        self, messages: list[dict[str, Any]], last: dict[str, Any], tools: Any, is_thinking: bool
    ) -> tuple[str, str | None, list[dict[str, Any]] | None, str]:
        if tools:
            if last.get("role") == "user":
                query = str(last.get("content", ""))[:128] or "検索"
                tool_calls = [
                    {
                        "id": "call_mock_1",
                        "type": "function",
                        "function": {
                            "name": "web_search",
                            "arguments": json.dumps({"query": query, "count": 3}, ensure_ascii=False),
                        },
                    }
                ]
                return "", "ユーザーの質問に答えるため、まず web_search で情報を集めます。", tool_calls, "tool_calls"

            if last.get("role") == "tool":
                tool_result_text = str(last.get("content", ""))
                url_match = re.search(r"https?://\S+", tool_result_text)
                already_fetched = any(
                    m.get("role") == "assistant"
                    and any(tc.get("function", {}).get("name") == "web_fetch" for tc in (m.get("tool_calls") or []))
                    for m in messages
                )
                if url_match and not already_fetched:
                    url = url_match.group(0).rstrip('",}')
                    tool_calls = [
                        {
                            "id": "call_mock_2",
                            "type": "function",
                            "function": {
                                "name": "web_fetch",
                                "arguments": json.dumps({"url": url}, ensure_ascii=False),
                            },
                        }
                    ]
                    return "", "検索結果の中から有力なページを 1 件取得します。", tool_calls, "tool_calls"

                # ツール結果をもとに最終回答を組み立てる。
                titles = re.findall(r'"title":\s*"([^"]*)"', tool_result_text)
                count_desc = f"{len(titles)} 件" if titles else "情報"
                title_desc = f"（例: {titles[0]}）" if titles else ""
                answer = f"収集した{count_desc}の資料{title_desc}をもとに回答します（モック応答）。"
                return answer, "収集した資料を踏まえて最終回答をまとめます。", None, "stop"

            # tools はあるが上記に当てはまらない初回メッセージ等 → 検索から開始。
            query = str(last.get("content", "質問"))[:128]
            tool_calls = [
                {
                    "id": "call_mock_1",
                    "type": "function",
                    "function": {
                        "name": "web_search",
                        "arguments": json.dumps({"query": query, "count": 3}, ensure_ascii=False),
                    },
                }
            ]
            return "", "まず検索します。", tool_calls, "tool_calls"

        # tools なし: 最後の user 発言を引用した単純な回答。
        user_text = str(last.get("content", ""))
        has_citation_instruction = any(
            m.get("role") == "system" and "出典" in str(m.get("content", "")) for m in messages
        )
        reasoning = f"「{user_text[:40]}」について考えます。" if is_thinking else None
        if has_citation_instruction:
            answer = f"「{user_text}」についてお答えします。資料に基づく回答です。[1]"
        else:
            answer = f"「{user_text}」についてお答えします（モック応答）。"
        return answer, reasoning, None, "stop"


def _error_body_dict(message: str) -> dict[str, Any]:
    # 公式ガイド（Search & Fetch）の Errors: OpenAI 互換のエラー形式。param は常に null。
    return {"error": {"message": message, "type": "invalid_request_error", "param": None, "code": "invalid_request_error"}}


class MockServer(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.requests: list[dict[str, Any]] = []
        # mock:429 / mock:408 は「同じクエリ文字列」ごとに最初の 2 回だけ失敗させ、3 回目で成功させる。
        # サーバーインスタンスごとに持つので、モックを再起動すれば回数はリセットされる。
        self.search_429_counts: dict[str, int] = {}
        self.search_408_counts: dict[str, int] = {}


def make_server(port: int = 8787) -> tuple[MockServer, int]:
    """モックサーバーを作成する（まだ起動しない）。port=0 で空きポートを自動選択。"""
    server = MockServer(("127.0.0.1", port), MockHandler)
    return server, server.server_address[1]


def serve_in_thread(server: MockServer) -> threading.Thread:
    """バックグラウンドスレッドでサーバーを動かす。終了時は server.shutdown() を呼ぶこと。"""
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return thread


def main() -> None:
    parser = argparse.ArgumentParser(description="DeepSeek API のローカルモックサーバー")
    parser.add_argument("--port", type=int, default=8787)
    args = parser.parse_args()

    server, port = make_server(args.port)
    print(f"モックサーバーを起動しました: http://127.0.0.1:{port}")
    print("Ctrl+C で終了します。")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.shutdown()
        server.server_close()


if __name__ == "__main__":
    main()