#!/usr/bin/env python3
"""验证全文搜索依赖；只操作随机命名的临时索引，不触碰业务文件。"""
import io
import json
import os
import pathlib
import subprocess
import sys
import time
import urllib.error
import urllib.request
import uuid
import zipfile

RUNTIME = pathlib.Path(os.environ.get('FTS_RUNTIME', '/opt/DockerDate/cloudreve-fts'))
STATE = RUNTIME / '.verification-index'
MEILI = 'http://127.0.0.1:7700'
TIKA = 'http://127.0.0.1:9998'
KEY = next(line.split('=', 1)[1].strip() for line in (RUNTIME / '.env').read_text().splitlines() if line.startswith('MEILI_MASTER_KEY='))


def request(base, path, method='GET', payload=None, auth=True, content_type='application/json'):
    headers = {'Accept': 'text/plain' if base == TIKA and path in ('/version', '/tika') else 'application/json'}
    if auth and base == MEILI:
        headers['Authorization'] = 'Bearer ' + KEY
    if payload is not None:
        headers['Content-Type'] = content_type
        if not isinstance(payload, bytes):
            payload = json.dumps(payload, ensure_ascii=False).encode()
    req = urllib.request.Request(base + path, data=payload, headers=headers, method=method)
    with urllib.request.urlopen(req, timeout=45) as response:
        data = response.read()
        try:
            return json.loads(data)
        except (ValueError, UnicodeDecodeError):
            return data.decode('utf-8')


def wait_task(task):
    deadline = time.monotonic() + 60
    while time.monotonic() < deadline:
        result = request(MEILI, '/tasks/' + str(task['taskUid']))
        if result['status'] == 'succeeded':
            return
        if result['status'] in ('failed', 'canceled'):
            raise RuntimeError('索引任务失败：' + str(result.get('error')))
        time.sleep(0.2)
    raise TimeoutError('等待索引任务超时')


def verify_runtime():
    for name in ('cloudreve-fts-tika', 'cloudreve-fts-meilisearch'):
        result = json.loads(subprocess.check_output(['docker', 'inspect', name]))[0]
        host = result['HostConfig']
        assert result['State']['Running'] and not result['State']['OOMKilled'], '容器未运行或曾被内存限制终止'
        assert host['CpusetCpus'] == '0-3'
        assert host['Memory'] == 2 * 1024**3
        assert host['ReadonlyRootfs']
        assert host['RestartPolicy']['Name'] == 'unless-stopped'
        assert host['CapDrop'] == ['ALL']
        for mappings in result['NetworkSettings']['Ports'].values():
            assert mappings and all(item['HostIp'] == '127.0.0.1' for item in mappings), '端口未限制为本机'
        image = json.loads(subprocess.check_output(['docker', 'image', 'inspect', result['Image']]))[0]
        assert image['Architecture'] == 'arm64'
        processes = subprocess.check_output(['docker', 'top', name, '-eo', 'pid'], text=True).splitlines()[1:]
        count = 0
        for line in processes:
            pid = line.strip()
            for status in pathlib.Path('/proc', pid, 'task').glob('*/status'):
                affinity = next(x.split(':', 1)[1].strip() for x in status.read_text().splitlines() if x.startswith('Cpus_allowed_list:'))
                assert affinity == '0-3', '线程亲和力不正确：' + affinity
                count += 1
        assert count > 0
        print(name + '：ARM64、回环端口、安全配置、所有线程0–3通过（' + str(count) + '线程）')
    assert (RUNTIME / '.env').stat().st_mode & 0o777 == 0o600
    assert pathlib.Path(subprocess.check_output(['findmnt', '-T', str(RUNTIME), '-no', 'TARGET'], text=True).strip()) == pathlib.Path('/')
    print('私密配置600权限及系统盘存储通过')


def make_docx(text):
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, 'w') as archive:
        archive.writestr('[Content_Types].xml', '<?xml version="1.0"?><Types xmlns="http://schemas.openxmlformats.org/package/2006/content-types"><Default Extension="rels" ContentType="application/vnd.openxmlformats-package.relationships+xml"/><Default Extension="xml" ContentType="application/xml"/><Override PartName="/word/document.xml" ContentType="application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"/></Types>')
        archive.writestr('_rels/.rels', '<?xml version="1.0"?><Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships"><Relationship Id="rId1" Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument" Target="word/document.xml"/></Relationships>')
        archive.writestr('word/document.xml', '<?xml version="1.0" encoding="UTF-8"?><w:document xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main"><w:body><w:p><w:r><w:t>' + text + '</w:t></w:r></w:p></w:body></w:document>')
    return buffer.getvalue()


def verify_search(uid):
    hits = request(MEILI, '/indexes/' + uid + '/search', 'POST', {'q': '蓝鲸', 'filter': 'owner_id = 1'})['hits']
    assert any(hit['id'] == 1 for hit in hits), '中文搜索未命中'
    print('中文正文检索及用户过滤通过')


def main():
    mode = sys.argv[1] if len(sys.argv) > 1 else '--test'
    verify_runtime()
    assert request(MEILI, '/health')['status'] == 'available'
    assert '3.2.3' in request(TIKA, '/version')
    assert request(MEILI, '/version')['pkgVersion'] == '1.54.3'
    print('两服务健康及版本通过')
    try:
        request(MEILI, '/indexes', auth=False)
        raise AssertionError('无密钥请求被错误放行')
    except urllib.error.HTTPError as error:
        assert error.code in (401, 403)
        print('Meilisearch拒绝未认证请求通过')
    if mode == '--finish':
        uid = STATE.read_text().strip()
        assert uid.startswith('fts_verify_') and uid.removeprefix('fts_verify_').isalnum()
        verify_search(uid)
        wait_task(request(MEILI, '/indexes/' + uid, 'DELETE'))
        STATE.unlink()
        print('重启后索引持久化通过；临时测试索引已删除')
        return
    assert mode in ('--test', '--keep'), '未知验证模式'
    assert not STATE.exists(), '上一次临时索引尚未清理，请运行--finish'
    text = '蓝鲸计划用于验证中文全文检索。Cloudreve dependency smoke test.'
    extracted = request(TIKA, '/tika', 'PUT', text.encode(), auth=False, content_type='text/plain; charset=UTF-8')
    assert '蓝鲸' in extracted
    print('UTF-8中文TXT提取通过')
    metadata = request(TIKA, '/rmeta/text', 'PUT', make_docx(text), auth=False, content_type='application/vnd.openxmlformats-officedocument.wordprocessingml.document')
    content = metadata[0]['X-TIKA:content']
    assert '蓝鲸' in content
    print('DOCX中文正文提取通过')
    uid = 'fts_verify_' + uuid.uuid4().hex
    created = False
    keep = False
    try:
        wait_task(request(MEILI, '/indexes', 'POST', {'uid': uid, 'primaryKey': 'id'}))
        created = True
        wait_task(request(MEILI, '/indexes/' + uid + '/settings', 'PATCH', {'searchableAttributes': ['content'], 'filterableAttributes': ['owner_id']}))
        wait_task(request(MEILI, '/indexes/' + uid + '/documents', 'POST', [{'id': 1, 'owner_id': 1, 'content': content}]))
        verify_search(uid)
        if mode == '--keep':
            STATE.write_text(uid)
            STATE.chmod(0o600)
            keep = True
            print('保留临时样本用于重启持久化验证')
    finally:
        if created and not keep:
            wait_task(request(MEILI, '/indexes/' + uid, 'DELETE'))
            print('临时测试索引已删除')


if __name__ == '__main__':
    main()