import hashlib
import hmac
import json
import os
import sqlite3
import struct
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from wechat_tool.cipher import (PAGE, RESERVE, aes_cbc, decrypt_database,
                                hmac_key, matches, wal_checksum)
from wechat_tool.mac_connect import MacConnection, parse_keys
from wechat_tool.wechat_native import message_rows
from wechat_tool import store


def sqlite_fixture(text):
    db = sqlite3.connect(':memory:')
    db.execute('PRAGMA page_size=4096')
    db.execute('CREATE TABLE t (value TEXT)')
    db.execute('INSERT INTO t VALUES (?)', (text,))
    data = db.serialize()
    db.close()
    return reserved_pages(data)


def reserved_pages(data):
    pages = []
    # Repack these small, single-leaf fixture pages with SQLCipher's reserved tail.
    for index in range(0, len(data), PAGE):
        page = bytearray(data[index:index+PAGE])
        offset = 100 if index == 0 else 0
        assert page[offset] == 13  # no overflow, free blocks or internal nodes
        count = int.from_bytes(page[offset+3:offset+5], 'big')
        start = int.from_bytes(page[offset+5:offset+7], 'big') or PAGE
        cells = bytes(page[start:])
        page[start-RESERVE:PAGE-RESERVE] = cells
        page[PAGE-RESERVE:] = bytes(RESERVE)
        page[offset+5:offset+7] = (start-RESERVE).to_bytes(2, 'big')
        for i in range(count):
            pos = offset+8+2*i
            value = int.from_bytes(page[pos:pos+2], 'big')
            page[pos:pos+2] = (value-RESERVE).to_bytes(2, 'big')
        if not index:
            page[20] = RESERVE
        pages.append(bytes(page))
    return pages


def encrypted_page(plain, number, key, salt):
    start = 16 if number == 1 else 0
    iv = os.urandom(16)
    ciphertext = aes_cbc(plain[start:-RESERVE], key, iv, True)
    mac = hmac.new(hmac_key(key, salt), ciphertext+iv+struct.pack('<I', number), hashlib.sha512).digest()
    return (salt if number == 1 else b'') + ciphertext + iv + mac


def native_fixture(root, invalid=False):
    (root/'contact').mkdir(parents=True)
    (root/'message').mkdir()
    group = 'test-group@chatroom'
    db = sqlite3.connect(root/'contact/contact.db')
    db.execute('CREATE TABLE contact (username TEXT, nick_name TEXT, remark TEXT)')
    db.executemany('INSERT INTO contact VALUES (?,?,?)', [(group,'测试群',''),('sender-a','张工','张工程师')])
    db.commit()
    db.close()
    table = 'Msg_'+hashlib.md5(group.encode()).hexdigest()
    db = sqlite3.connect(root/'message/message_0.db')
    db.execute('CREATE TABLE Name2Id (user_name TEXT)')
    db.executemany('INSERT INTO Name2Id VALUES (?)', [(group,),('sender-a',)])
    db.execute('CREATE TABLE '+table+' (local_id INTEGER, server_id INTEGER, create_time INTEGER, message_content, local_type INTEGER, real_sender_id INTEGER)')
    db.execute('INSERT INTO '+table+' VALUES (1,1001,1790816400,?,1,2)', ('测试比例怎么确定？',))
    if invalid:
        db.execute('INSERT INTO '+table+' VALUES (2,1002,1790816401,?,1,999)', ('不能错误归属',))
    db.commit()
    db.close()


class CipherTests(unittest.TestCase):
    def test_native_aes_nist_vector(self):
        key = bytes.fromhex('603deb1015ca71be2b73aef0857d77811f352c073b6108d72d9810a30914dff4')
        iv = bytes.fromhex('000102030405060708090a0b0c0d0e0f')
        plain = bytes.fromhex('6bc1bee22e409f96e93d7e117393172a')
        expected = bytes.fromhex('f58c4c04d6e5f1ba779eabfb5f7bfbd6')
        self.assertEqual(aes_cbc(plain,key,iv,True), expected)
        self.assertEqual(aes_cbc(expected,key,iv), plain)

    def test_authenticated_decryption_and_corruption(self):
        key, salt = os.urandom(32), os.urandom(16)
        pages = sqlite_fixture('正确内容')
        data = b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(pages,1))
        self.assertTrue(matches(data[:PAGE], key))
        self.assertFalse(matches(data[:PAGE], os.urandom(32)))
        with tempfile.TemporaryDirectory() as directory:
            src, dst = Path(directory)/'source.db', Path(directory)/'readable.db'
            src.write_bytes(data)
            decrypt_database(src,dst,key)
            db=sqlite3.connect(dst)
            self.assertEqual(db.execute('SELECT value FROM t').fetchone()[0], '正确内容')
            db.close()
            self.assertEqual(src.read_bytes(),data)
            broken=bytearray(data);broken[PAGE+30]^=1;src.write_bytes(broken)
            other=Path(directory)/'broken.db'
            with self.assertRaises(ValueError):
                decrypt_database(src,other,key)
            self.assertFalse(other.exists())
            with self.assertRaises(FileExistsError):
                src.write_bytes(data)
                decrypt_database(src,dst,key)
            self.assertTrue(dst.exists())

    def test_wal_replays_only_commits_and_rejects_bad_checksum(self):
        key,salt=os.urandom(32),os.urandom(16)
        before,after,last=sqlite_fixture('before'),sqlite_fixture('committed'),sqlite_fixture('uncommitted')
        main=b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(before,1))
        header=struct.pack('>IIIIII',0x377f0682,3007000,PAGE,0,123,456)
        checksum=wal_checksum(header)
        wal=bytearray(header+struct.pack('>II',*checksum))
        for number,page in list(enumerate(after,1))+[(2,last[1])]:
            count=len(after) if len(wal)==32+(len(after)-1)*(PAGE+24) else 0
            payload=encrypted_page(page,number,key,salt)
            first=struct.pack('>II',number,count)
            checksum=wal_checksum(first+payload,checksum)
            wal.extend(first+struct.pack('>II',123,456)+struct.pack('>II',*checksum)+payload)
        with tempfile.TemporaryDirectory() as directory:
            src,dst=Path(directory)/'source.db',Path(directory)/'out.db'
            src.write_bytes(main);walpath=Path(str(src)+'-wal');walpath.write_bytes(wal)
            result=decrypt_database(src,dst,key)
            self.assertEqual(result['wal_frames'],len(after))
            db=sqlite3.connect(dst)
            self.assertEqual(db.execute('SELECT value FROM t').fetchone()[0],'committed')
            db.close()
            wal[-1]^=1;walpath.write_bytes(wal)
            with self.assertRaises(ValueError):
                decrypt_database(src,Path(directory)/'bad.db',key)


class NativeTests(unittest.TestCase):
    def test_encrypted_native_pipeline(self):
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory)/'source'/'account'/'db_storage'
            native_fixture(root)
            keys, originals = [], {}
            for path in MacConnection.files(root):
                key,salt=os.urandom(32),os.urandom(16)
                keys.append(key)
                pages=reserved_pages(path.read_bytes())
                data=b''.join(encrypted_page(p,i,key,salt) for i,p in enumerate(pages,1))
                path.write_bytes(data)
                originals[path]=data
            service=MacConnection(Path(directory)/'workspace');service.job={'status':'running'}
            with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
                result=service.import_account(root,keys,False)
            self.assertEqual(result['added'],1)
            for path,data in originals.items():
                self.assertEqual(path.read_bytes(),data)
            db=store.connect(service.data)
            self.assertEqual(db.execute('SELECT sender FROM messages').fetchone()[0],'张工程师')
            db.close()

    def test_key_json_variants_never_return_raw_keys(self):
        key='ab'*32
        self.assertEqual(parse_keys({'a':{'key_hex':key},'b':{'key':"x'"+key+'cd'*16+"'"}}),[bytes.fromhex(key)])
        with self.assertRaises(ValueError):
            parse_keys({'password':'微信密码不是密钥'})
        with tempfile.TemporaryDirectory() as directory:
            service=MacConnection(directory)
            state=service.load_keys({'key':key})
            self.assertTrue(state['keys_ready'])
            self.assertNotIn(key,json.dumps(state))
            self.assertEqual(list(Path(directory).rglob('*')),[])

    def test_native_mapping_and_atomic_idempotent_import(self):
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory)/'source'/'account'/'db_storage'
            native_fixture(root)
            rows=list(message_rows(root,'account'))
            self.assertEqual(rows[0]['sender'],'张工程师')
            self.assertEqual(rows[0]['kind'],'group')
            self.assertEqual(rows[0]['conversation'],'测试群')
            service=MacConnection(Path(directory)/'workspace')
            service.job={'status':'running'}
            with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
                one=service.import_account(root,[],False)
                two=service.import_account(root,[],False)
            self.assertEqual((one['added'],two['duplicates']),(1,1))
            db=store.connect(service.data)
            self.assertEqual(db.execute('SELECT COUNT(*) FROM messages').fetchone()[0],1)
            db.close()

    def test_unknown_identity_aborts_entire_import(self):
        with tempfile.TemporaryDirectory() as directory:
            root=Path(directory)/'source'/'account'/'db_storage'
            native_fixture(root,True)
            service=MacConnection(Path(directory)/'workspace');service.job={'status':'running'}
            with patch('wechat_tool.mac_connect.wechat_pid',return_value=None):
                with self.assertRaises(ValueError):
                    service.import_account(root,[],False)
            db=store.connect(service.data)
            self.assertEqual(db.execute('SELECT COUNT(*) FROM messages').fetchone()[0],0)
            db.close()

    def test_running_wechat_and_non_mac_are_actionable(self):
        with tempfile.TemporaryDirectory() as directory:
            service=MacConnection(directory)
            with patch('wechat_tool.mac_connect.wechat_pid',return_value=123):
                with self.assertRaisesRegex(ValueError,'退出微信'):
                    service.import_account(Path(directory),[],False)
            with patch('wechat_tool.mac_connect.platform.system',return_value='Windows'):
                with self.assertRaisesRegex(ValueError,'目标 Mac'):
                    service.scan()