import base64
import json
import logging
from rpa_engine.egress_channels import source_bound_requests_session
from .dy_util import (
    trans_cookies,
    generate_msToken,
    generate_a_bogus,
    splice_url,
    generate_webid,
    normalize_client_cert,
    resolve_proto_device_id,
    DEFAULT_USER_AGENT,
)

logger = logging.getLogger("douyin_im.auth")


def _parse_storage_json(raw) -> dict | None:
    """Parse localStorage JSON (supports nested data wrapper)."""
    if not raw:
        return None
    if isinstance(raw, dict):
        data = raw
    else:
        text = str(raw).strip()
        data = None
        for _ in range(4):
            try:
                parsed = json.loads(text)
            except Exception:
                break
            if isinstance(parsed, str):
                text = parsed
                continue
            if isinstance(parsed, dict):
                data = parsed
                break
            break
        if not data:
            return None

    inner = data.get("data")
    if isinstance(inner, str):
        try:
            inner = json.loads(inner)
        except Exception:
            inner = None
    if isinstance(inner, dict):
        return inner
    return data


class DouyinAuth:
    def __init__(self):
        self.cookie = None
        self.cookie_str = None
        self.private_key = None
        self.ticket = None
        self.ts_sign = None
        self.client_cert = None
        self.ree_public_key = None
        self.uid = None
        self.msToken = None
        self.web_id = None
        self.source_ip = ""
        self.user_agent = None

    def perepare_auth(
        self,
        cookieStr: str,
        web_protect_: str = "",
        keys_: str = "",
        user_agent: str = "",
    ):
        self.cookie = trans_cookies(cookieStr)
        self.cookie_str = cookieStr
        self.msToken = self.cookie["msToken"] if "msToken" in self.cookie else generate_msToken()
        self.cookie["msToken"] = self.msToken
        self.cookie_str = "; ".join([f"{k}={v}" for k, v in self.cookie.items()])
        web_data = _parse_storage_json(web_protect_)
        if web_data:
            try:
                self.ticket = web_data.get("ticket") or ""
                self.ts_sign = web_data.get("ts_sign") or ""
                self.client_cert = web_data.get("client_cert") or ""
            except Exception as e:
                logger.debug(f"web_protect parse failed: {e}")

        keys_data = _parse_storage_json(keys_)
        if keys_data:
            try:
                self.private_key = keys_data.get("ec_privateKey") or keys_data.get("privateKey") or ""
                if self.private_key:
                    self.ree_public_key = base64.b64encode(self.private_key.encode()).decode()
            except Exception as e:
                logger.debug(f"keys parse failed: {e}")

        if user_agent:
            # 让签名上下文记住调用方 UA：query_my_uid / generate_webid 等后续
            # 请求会复用它，避免退回硬编码 DEFAULT_USER_AGENT 造成 UA 不一致。
            self.user_agent = user_agent

    def is_sign_ready(self) -> bool:
        return bool(
            self.private_key
            and self.ticket
            and self.ts_sign
            and self.client_cert
        )

    @classmethod
    def from_im_session(cls, session) -> "DouyinAuth":
        """从 DouyinImSession 构建 HTTP 发送/拉会话用的签名上下文。

        注意：frontier WS 的 sdk_cert/ts_sign 只用于长连接，不能覆盖 web_protect，
        否则 bd-ticket-guard 与 protobuf 签名会与 ticket 失配 -> 7911。
        """
        auth = cls()
        auth.perepare_auth(
            session.cookie_header(),
            session.web_protect_str,
            session.keys_str,
            user_agent=session.user_agent or DEFAULT_USER_AGENT,
        )
        auth.web_id = session.web_id or session.device_id or None
        # device_id 是设备注册号（query/user 的 id），不是账号 UID；
        # my_uid 只作为最后兜底，由 resolve_proto_device_id 内部处理。
        auth.device_id = resolve_proto_device_id(
            session.device_id, session.web_id, session.my_uid
        )
        auth.source_ip = str(getattr(session, "egress_source_ip", "") or "")
        # web_protect 缺 client_cert 时，才用 frontier 抓包证书兜底（不覆盖 ts_sign）
        if not auth.client_cert and getattr(session, "sdk_cert", ""):
            auth.client_cert = normalize_client_cert(session.sdk_cert)
        elif auth.client_cert:
            auth.client_cert = normalize_client_cert(auth.client_cert)
        return auth

    def get_uid(self):
        if self.uid is None:
            # 优先从 cookie 尝试提取，否则请求接口
            for k in ("uid_tt", "uid_tt_ss"):
                if self.cookie and self.cookie.get(k):
                    try:
                        self.uid = int(self.cookie.get(k))
                        return self.uid
                    except ValueError:
                        pass
            try:
                self.uid = self.query_my_uid()
            except Exception:
                pass
        return self.uid

    def query_my_uid(self) -> int:
        ua = self.user_agent or DEFAULT_USER_AGENT
        url = 'https://www.douyin.com/aweme/v1/web/query/user/'
        headers = {
            "User-Agent": ua,
            "Referer": "https://www.douyin.com/",
            "Accept": "application/json, text/plain, */*",
        }
        params = {
            "device_platform": "webapp",
            "aid": "6383",
            "channel": "channel_pc_web",
            "publish_video_strategy_type": "2",
            "verifyFp": self.cookie.get('s_v_web_id', ''),
            "fp": self.cookie.get('s_v_web_id', ''),
            "webid": generate_webid(self, "https://www.douyin.com/"),
            "msToken": self.msToken
        }
        query = splice_url(params)
        abogus = generate_a_bogus(query, user_agent=ua)
        params['a_bogus'] = abogus
        
        with source_bound_requests_session(self.source_ip) as client:
            resp = client.get(
                url,
                params=params,
                headers=headers,
                cookies=self.cookie,
                verify=False,
                timeout=10,
            )
        resp_json = resp.json()
        return int(resp_json['user_uid'])