import contextlib
import json
import os
import shutil
import sys
import traceback
import uuid
from pathlib import Path

from domain import auto_match, enrich, extract
from exports import export_expense, export_ppt, export_travel

PROTOCOL = sys.stdout
SUPPORTED = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif', '.pdf'}


def emit(value):
    PROTOCOL.write(json.dumps(value, ensure_ascii=False) + '\n')
    PROTOCOL.flush()


def scan(request):
    import onnxruntime
    import pypdfium2 as pdfium
    from PIL import Image
    from rapidocr_onnxruntime import RapidOCR
    root = Path(request['path'])
    destination = Path(request['workspacePath'])
    if not root.is_dir():
        raise ValueError('请选择有效的报账材料文件夹')
    destination.mkdir(parents=True, exist_ok=True)
    state = dict(rootPath=request.get('folderName', root.name), invoices=[], payments=[], photos=[], matches=[], warnings=[], directoryPaymentTotal=0)
    files = []
    for source in sorted(root.rglob('*')):
        if source.is_file() and source.suffix.lower() in SUPPORTED:
            parts = source.relative_to(root).parts[:-1]
            kind = next((kind for label, kind in [('发票', 'invoice'), ('付款截图', 'payment'), ('实物照片', 'photo')] if label in parts), None)
            if kind:
                files.append((source, kind))
    if not files:
        raise ValueError('目录中没有可处理材料，请确认包含“发票”“付款截图”子文件夹')
    emit(dict(event='progress', progress=0, message='正在加载本地 RapidOCR 模型'))
    onnxruntime.disable_telemetry_events()
    recognizer = RapidOCR()
    for index, (source, kind) in enumerate(files):
        identifier = str(uuid.uuid4())
        stored = destination / (identifier + source.suffix.lower())
        shutil.copyfile(source, stored)
        item = dict(id=identifier, name=source.name, path=str(stored), type=kind, size=stored.stat().st_size, matched=False,
                    previewPath='', ocr=extract([]))
        item['ocr']['status'] = 'pending'
        try:
            if source.suffix.lower() == '.pdf':
                with contextlib.closing(pdfium.PdfDocument(str(stored))) as document:
                    if not len(document):
                        raise ValueError('PDF 没有可预览页面')
                    with contextlib.closing(document[0]) as page:
                        with contextlib.closing(page.render(scale=140 / 72)) as bitmap:
                            image = bitmap.to_pil().convert('RGB')
            else:
                with Image.open(stored) as original:
                    image = original.convert('RGB')
            preview = destination / (identifier + '-preview.png')
            image.save(preview, format='PNG')
            item['previewPath'] = str(preview)
            if kind != 'photo':
                result, _ = recognizer(str(preview if source.suffix.lower() == '.pdf' else stored))
                item['ocr'] = extract([str(entry[1]).strip() for entry in result or [] if len(entry) >= 3 and str(entry[1]).strip()])
        except Exception as error:
            item['ocr']['status'] = 'failed'
            item['ocr']['message'] = str(error)
            state['warnings'].append(source.name + '：识别或预览失败，可人工配对')
        state[{'invoice': 'invoices', 'payment': 'payments', 'photo': 'photos'}[kind]].append(item)
        emit(dict(event='progress', progress=(index + 1) / len(files) * .95, message=f'已处理 {index + 1}/{len(files)} · {source.name}'))
    for field, label in [('invoices', '发票'), ('payments', '付款截图'), ('photos', '实物照片')]:
        if not state[field]:
            state['warnings'].append('未找到“' + label + '”文件夹或支持的文件')
    emit(dict(event='progress', progress=.97, message='正在执行自动配对与分类'))
    return enrich(auto_match(state))


def dispatch(request):
    operation = request['operation']
    if operation == 'scan':
        return scan(request)
    state = request['state']
    if operation == 'refresh':
        return enrich(state)
    if not state['matches']:
        raise ValueError('至少完成一组核对后才能导出')
    destination = Path(request['destination'])
    temporary = destination.with_name('.' + str(uuid.uuid4()) + destination.suffix)
    try:
        if operation == 'ppt':
            export_ppt(state, temporary, request.get('classified', True))
        elif operation == 'travel':
            export_travel(state, temporary)
        elif operation == 'expense':
            template = Path(__file__).parent / 'personal-expense-template.xlsx'
            export_expense(state, temporary, template, request.get('purposes', {}), request.get('signatures', []))
        else:
            raise ValueError('未知操作：' + operation)
        os.replace(temporary, destination)
    finally:
        temporary.unlink(missing_ok=True)
    return dict(destination=str(destination))


if __name__ == '__main__':
    try:
        request = json.loads(sys.stdin.readline())
        with contextlib.redirect_stdout(sys.stderr):
            result = dispatch(request)
        emit(dict(event='result', result=result))
    except Exception as error:
        traceback.print_exc(file=sys.stderr)
        emit(dict(event='error', message=str(error)))
        sys.exit(1)