import os
import base64
import json
import hashlib
import logging
import re
import sqlite3
import sys
import unittest
import zipfile
import importlib
from pathlib import Path
from tempfile import TemporaryDirectory


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))


class TestChatExportHtmlPaging(unittest.TestCase):
    @staticmethod
    def _decode_runtime_js(script: str) -> str:
        match = re.search(r"const _0=(\[.*?\]),_1=(\[.*?\]);let", script)
        if match is None:
            raise AssertionError("unable to decode native HTML export runtime")
        chunks = json.loads(match.group(1))
        key = json.loads(match.group(2))
        packed = base64.b64decode("".join(chunks))
        return bytes(value ^ key[index % len(key)] for index, value in enumerate(packed)).decode("utf-8")

    def _reload_export_modules(self):
        import wechat_decrypt_tool.app_paths as app_paths
        import wechat_decrypt_tool.chat_helpers as chat_helpers
        import wechat_decrypt_tool.media_helpers as media_helpers
        import wechat_decrypt_tool.chat_export_service as chat_export_service

        importlib.reload(app_paths)
        importlib.reload(chat_helpers)
        importlib.reload(media_helpers)
        importlib.reload(chat_export_service)
        return chat_export_service

    def _seed_contact_db(self, path: Path, *, account: str, username: str) -> None:
        conn = sqlite3.connect(str(path))
        try:
            conn.execute(
                """
                CREATE TABLE contact (
                    username TEXT,
                    remark TEXT,
                    nick_name TEXT,
                    alias TEXT,
                    local_type INTEGER,
                    verify_flag INTEGER,
                    big_head_url TEXT,
                    small_head_url TEXT
                )
                """
            )
            conn.execute(
                """
                CREATE TABLE stranger (
                    username TEXT,
                    remark TEXT,
                    nick_name TEXT,
                    alias TEXT,
                    local_type INTEGER,
                    verify_flag INTEGER,
                    big_head_url TEXT,
                    small_head_url TEXT
                )
                """
            )
            conn.execute(
                "INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (account, "", "Me", "", 1, 0, "", ""),
            )
            conn.execute(
                "INSERT INTO contact VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (username, "", "Friend", "", 1, 0, "", ""),
            )
            conn.commit()
        finally:
            conn.close()

    def _seed_session_db(self, path: Path, *, username: str) -> None:
        conn = sqlite3.connect(str(path))
        try:
            conn.execute(
                """
                CREATE TABLE SessionTable (
                    username TEXT,
                    is_hidden INTEGER,
                    sort_timestamp INTEGER
                )
                """
            )
            conn.execute(
                "INSERT INTO SessionTable VALUES (?, ?, ?)",
                (username, 0, 1735689600),
            )
            conn.commit()
        finally:
            conn.close()

    def _seed_message_db(self, path: Path, *, account: str, username: str, total: int) -> None:
        conn = sqlite3.connect(str(path))
        try:
            conn.execute("CREATE TABLE Name2Id (rowid INTEGER PRIMARY KEY, user_name TEXT)")
            conn.execute("INSERT INTO Name2Id(rowid, user_name) VALUES (?, ?)", (1, account))
            conn.execute("INSERT INTO Name2Id(rowid, user_name) VALUES (?, ?)", (2, username))

            table_name = f"msg_{hashlib.md5(username.encode('utf-8')).hexdigest()}"
            conn.execute(
                f"""
                CREATE TABLE {table_name} (
                    local_id INTEGER,
                    server_id INTEGER,
                    local_type INTEGER,
                    sort_seq INTEGER,
                    real_sender_id INTEGER,
                    create_time INTEGER,
                    message_content TEXT,
                    compress_content BLOB
                )
                """
            )

            # Generate lots of plain text messages with unique markers.
            rows = []
            base_ts = 1735689600
            for i in range(1, total + 1):
                marker = f"MSG{i:04d}"
                real_sender_id = 1 if (i % 2 == 0) else 2
                rows.append((i, 100000 + i, 1, i, real_sender_id, base_ts + i, marker, None))

            conn.executemany(
                f"INSERT INTO {table_name} (local_id, server_id, local_type, sort_seq, real_sender_id, create_time, message_content, compress_content) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                rows,
            )
            conn.commit()
        finally:
            conn.close()

    def _prepare_account(self, root: Path, *, account: str, username: str, total: int) -> Path:
        account_dir = root / "output" / "databases" / account
        account_dir.mkdir(parents=True, exist_ok=True)
        self._seed_contact_db(account_dir / "contact.db", account=account, username=username)
        self._seed_session_db(account_dir / "session.db", username=username)
        self._seed_message_db(account_dir / "message_0.db", account=account, username=username, total=total)
        return account_dir

    def _create_job(self, manager, *, account: str, username: str, html_page_size: int):
        job = manager.create_job(
            account=account,
            source="decrypted",
            scope="selected",
            usernames=[username],
            export_format="html",
            start_time=None,
            end_time=None,
            include_hidden=False,
            include_official=False,
            include_media=False,
            media_kinds=[],
            message_types=[],
            output_dir=None,
            allow_process_key_extract=False,
            download_remote_media=False,
            html_page_size=html_page_size,
            privacy_mode=False,
            file_name=None,
        )

        # Export is async (thread). Allow enough time for a few thousand messages + zip writes.
        for _ in range(600):
            latest = manager.get_job(job.export_id)
            if latest and latest.status in {"done", "error", "cancelled"}:
                return latest
            import time as _time

            _time.sleep(0.05)
        self.fail("export job did not finish in time")

    def test_html_export_paging_inlines_latest_page_only(self):
        with TemporaryDirectory() as td:
            root = Path(td)
            account = "wxid_test"
            username = "wxid_friend"

            total_messages = 2300
            page_size = 1000
            self._prepare_account(root, account=account, username=username, total=total_messages)

            prev_data = os.environ.get("WECHAT_TOOL_DATA_DIR")
            try:
                os.environ["WECHAT_TOOL_DATA_DIR"] = str(root)
                svc = self._reload_export_modules()
                job = self._create_job(
                    svc.CHAT_EXPORT_MANAGER,
                    account=account,
                    username=username,
                    html_page_size=page_size,
                )
                self.assertEqual(job.status, "done", msg=job.error)

                self.assertTrue(job.zip_path and job.zip_path.exists())
                with zipfile.ZipFile(job.zip_path, "r") as zf:
                    names = set(zf.namelist())

                    html_path = next((n for n in names if n.endswith("/messages.html")), "")
                    self.assertTrue(html_path, msg="missing messages.html")
                    html_text = zf.read(html_path).decode("utf-8", errors="ignore")

                    # Paging UI + meta should exist for multi-page exports.
                    self.assertIn('id="wcePageMeta"', html_text)
                    self.assertIn('id="wcePager"', html_text)
                    self.assertIn('id="wceMessageList"', html_text)
                    self.assertIn('id="wceLoadPrevBtn"', html_text)

                    # Latest page is inlined; earliest page should not be present in messages.html.
                    self.assertIn("MSG2300", html_text)
                    self.assertNotIn("MSG0001", html_text)

                    conv_dir = html_path.rsplit("/", 1)[0]
                    page1_js = f"{conv_dir}/pages/page-0001.js"
                    self.assertIn(page1_js, names)
                    page1_text = zf.read(page1_js).decode("utf-8", errors="ignore")
                    self.assertIn("MSG0001", page1_text)
            finally:
                logging.shutdown()
                if prev_data is None:
                    os.environ.pop("WECHAT_TOOL_DATA_DIR", None)
                else:
                    os.environ["WECHAT_TOOL_DATA_DIR"] = prev_data

    def test_paged_runtime_verifies_fragments_only_when_loaded(self):
        from wechat_decrypt_tool.export_integrity import load_wce_integrity_native

        native = load_wce_integrity_native()
        runtime = self._decode_runtime_js(str(native.runtime_js()))
        fragment = str(
            native.page_fragment_js(
                "export-test",
                "conversations/demo/pages/page-0001.js",
                1,
                '<div data-render-type="text">hello</div>',
            )
        )

        self.assertNotIn("precheckPromise = precheckAllPages()", runtime)
        self.assertNotIn("const precheckAllPages = async", runtime)
        self.assertIn("s.src = pageSrc(n)", runtime)
        self.assertIn("window.__WCE_PAGE_SEEN__", runtime)
        self.assertNotIn("__WCE_PAGE_SEEN__", fragment)
        self.assertNotIn("__WCE_PAGE_PRECHECK__", fragment)
        self.assertIn("__WCE_PAGE_LOADED__", fragment)
        self.assertIn("s.onload =", runtime)
        self.assertIn("pendingPageLoads.has(n)", runtime)
        self.assertIn("clearTimeout(pending.timer)", runtime)
