import json
import requests as req_lib
from .base import NO_PROXY


class DeepSeekClient:
    def __init__(self, token, cookie_id=""):
        self.host = "https://chat.deepseek.com"
        self.token = token if token.startswith("Bearer") else f"Bearer {token}"
        self.cookie_id = cookie_id
        self.headers = {
            "content-type": "application/json",
            "authorization": self.token,
            "origin": self.host,
            "x-app-version": "2.0.0",
            "x-client-platform": "web",
            "x-client-locale": "zh_CN",
            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/149.0.0.0 Safari/537.36",
        }

    def _get_pow_headers(self):
        try:
            from deps.deep_sign import get_ds_pow
            resp = req_lib.post(
                f"{self.host}/api/v0/chat/create_pow_challenge",
                headers=self.headers,
                json={"target_path": "/api/v0/chat/completion"},
                timeout=30, verify=False, proxies=NO_PROXY,
            )
            biz_data = resp.json()["data"]["biz_data"]
            pow_resp = get_ds_pow(authorization=self.token, target_path="/api/v0/chat/completion", biz_data=biz_data)
            return {"x-ds-pow-response": pow_resp}
        except Exception as e:
            print(f"    [WARN] PoW签名失败: {e}")
            return {}

    def create_session(self):
        resp = req_lib.post(
            f"{self.host}/api/v0/chat_session/create",
            headers=self.headers, json={}, timeout=30, verify=False, proxies=NO_PROXY,
        )
        data = resp.json()
        biz_data = data.get("data", {}).get("biz_data", {})
        session = biz_data.get("chat_session", biz_data)
        if not session.get("id"):
            raise Exception(f"create_session 失败: {json.dumps(data, ensure_ascii=False)[:200]}")
        return session

    def chat(self, question, thinking=False):
        session = self.create_session()
        sid = session["id"]
        headers = {**self.headers, **self._get_pow_headers()}
        headers["referer"] = f"{self.host}/a/chat/s/{sid}"

        payload = {
            "chat_session_id": sid,
            "parent_message_id": None,
            "model_type": "default",
            "prompt": question,
            "ref_file_ids": [],
            "thinking_enabled": thinking,
            "search_enabled": True,
            "action": None,
            "preempt": False,
        }

        resp = req_lib.post(
            f"{self.host}/api/v0/chat/completion",
            headers=headers, json=payload, stream=True, timeout=300, verify=False, proxies=NO_PROXY,
        )
        resp.raise_for_status()

        raw_text = ""
        for line in resp.iter_lines(decode_unicode=True):
            if line:
                raw_text += line + "\n"

        if "Authorization Failed" in raw_text or "user is muted" in raw_text:
            raise Exception(raw_text[:200])

        return raw_text, json.dumps(payload, ensure_ascii=False)


def get_ds_session_id(result):
    raw_req = result.get("_raw_req", "")
    try:
        req = json.loads(raw_req) if raw_req else {}
        return req.get("chat_session_id", "")
    except Exception:
        return ""