import base64
import unittest

from core.enc_data import ZT_OUTER_CONFIGS
from core.privacykit_encrypt import (
    PASSPORT_ACCOUNT_IMAGE_PREFIX,
    PASSPORT_WCFG_KEY,
    PRIVACYKIT_SDK_ID,
    WEAPONHI_IMG_INITIAL,
    diagnose_passport_account_image,
    java_gzip,
    parse_passport_account_image,
    privacykit_atlas_encrypt_raw,
    weaponhi_dd21_from_wcfg,
    weaponhi_b,
    weaponhi_vimg_upload,
)


class PrivacyKitEncryptTests(unittest.TestCase):
    def test_java_gzip_header_matches_app_samples(self):
        """Java GZIPOutputStream 样本头：mtime=0、XFL=0、OS=255。"""

        packed = java_gzip(b"hello")

        self.assertEqual(packed[:10].hex(), "1f8b08000000000000ff")
        self.assertGreater(len(packed), len(b"hello"))

    def test_privacykit_atlas_encrypt_raw_uses_privacykit_outer_head(self):
        raw = privacykit_atlas_encrypt_raw(b"abc", epoch_seconds=1785044982)

        self.assertTrue(raw.startswith(ZT_OUTER_CONFIGS[PRIVACYKIT_SDK_ID]["head8"]))
        info = parse_passport_account_image(PASSPORT_ACCOUNT_IMAGE_PREFIX + base64.b64encode(raw).decode("ascii"))
        self.assertEqual(info.format_kind, "zt_outer")
        self.assertEqual(info.sdk_id, PRIVACYKIT_SDK_ID)
        self.assertEqual(info.inner_magic_hex, "dec0adde")
        self.assertEqual(info.inner_header_size, 0x20)

    def test_weaponhi_b_and_vimg_upload_are_parseable(self):
        encrypted = weaponhi_b("hello", epoch_seconds=1785044982)
        vimg = weaponhi_vimg_upload("hello", epoch_seconds=1785044982)

        self.assertEqual(vimg, PASSPORT_ACCOUNT_IMAGE_PREFIX + encrypted)
        info = parse_passport_account_image(vimg)
        self.assertTrue(info.ok)
        self.assertEqual(info.format_kind, "zt_outer")
        self.assertFalse(info.has_ai_hash)

    def test_parse_final_passport_ticket_shape_from_app(self):
        """登录票据是本地 Engine.pr 的 VIMG + AI，不是上传态 ZT。"""

        raw = bytes.fromhex("203868e844b7900a75ef44ce0702671f") + b"ticket-body"
        body = base64.b64encode(raw).decode("ascii")
        value = f"{PASSPORT_ACCOUNT_IMAGE_PREFIX}{body}$AI_{'1' * 32}"

        info = parse_passport_account_image(value)

        self.assertTrue(info.ok)
        self.assertTrue(info.has_vimg_prefix)
        self.assertTrue(info.has_ai_hash)
        self.assertEqual(info.format_kind, "weapon_pr")
        self.assertEqual(info.raw_head_hex, raw[:16].hex())
        self.assertEqual(info.ai_hash, "1" * 32)
        self.assertFalse(info.ai_matches_md5_raw)
        self.assertFalse(info.ai_matches_md5_base64)

    def test_weaponhi_dd21_reads_wcfg_a_y_q_z_when_cache_is_initial(self):
        """dd(21)：img 不是 VIMG_ 时，从 wcfg['a_y_q_z'] 读取并缓存。"""

        raw = bytes.fromhex("203868e844b7900a75ef44ce0702671f") + b"ticket-body"
        value = f"{PASSPORT_ACCOUNT_IMAGE_PREFIX}{base64.b64encode(raw).decode('ascii')}$AI_{'2' * 32}"

        result = weaponhi_dd21_from_wcfg({PASSPORT_WCFG_KEY: value}, img_cache=WEAPONHI_IMG_INITIAL)

        self.assertTrue(result.ok)
        self.assertEqual(result.source, "wcfg.a_y_q_z")
        self.assertEqual(result.value, value)
        self.assertEqual(result.cache_after, value)
        self.assertTrue(result.has_wcfg_value)
        self.assertTrue(result.is_final_ticket)
        self.assertFalse(result.is_upload_vimg)
        self.assertEqual(result.diagnosis["format_kind"], "weapon_pr")

    def test_weaponhi_dd21_reuses_vimg_cache_before_wcfg(self):
        """dd(21)：img 已是 VIMG_ 时直接返回缓存，不读新 wcfg 值。"""

        cached = weaponhi_vimg_upload("cached", epoch_seconds=1785044982)
        newer = weaponhi_vimg_upload("newer", epoch_seconds=1785044983)

        result = weaponhi_dd21_from_wcfg({PASSPORT_WCFG_KEY: newer}, img_cache=cached)

        self.assertTrue(result.ok)
        self.assertEqual(result.source, "img_cache")
        self.assertEqual(result.value, cached)
        self.assertEqual(result.cache_after, cached)
        self.assertTrue(result.has_wcfg_value)
        self.assertFalse(result.is_final_ticket)
        self.assertTrue(result.is_upload_vimg)

    def test_weaponhi_dd21_returns_initial_cache_when_wcfg_missing(self):
        """dd(21)：wcfg 无票据时返回默认 img，不应误判为可用票据。"""

        result = weaponhi_dd21_from_wcfg({}, img_cache=WEAPONHI_IMG_INITIAL)

        self.assertFalse(result.ok)
        self.assertEqual(result.source, "default_img_cache")
        self.assertEqual(result.value, WEAPONHI_IMG_INITIAL)
        self.assertFalse(result.has_wcfg_value)
        self.assertFalse(result.is_final_ticket)
        self.assertFalse(result.is_upload_vimg)
        self.assertEqual(result.diagnosis["format_kind"], "invalid")

    def test_diagnose_invalid_value_returns_structured_error(self):
        diag = diagnose_passport_account_image("VIMG_not_base64!")

        self.assertFalse(diag["ok"])
        self.assertEqual(diag["format_kind"], "invalid")
        self.assertIn("Error", diag["error"])


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