#!/usr/bin/env python3
"""
audiveris_to_musescore.py

使用 Audiveris 进行五线谱 OCR 识别的完整流水线。

功能：
1. 调用 Audiveris 识别五线谱图片/PDF
2. 输出 .mxl 格式的 MusicXML 文件

Usage:
    python audiveris_to_musescore.py <input_image_or_pdf> [output_dir]

Examples:
    python audiveris_to_musescore.py "D:/scores/sheet.png" "D:/output"

    # 处理 PDF
    python audiveris_to_musescore.py "D:/scores/piece.pdf" "D:/output"

Dependencies:
    - Audiveris (https://github.com/Audiveris/audiveris)
    - Python 3.8+

Installation:
    # Windows
    winget install --id Audiveris.Audiveris --accept-package-agreements --accept-source-agreements

    # 或手动下载 MSI: https://github.com/Audiveris/audiveris/releases
"""

import sys
import os
import subprocess
import tempfile
import shutil
import zipfile
import argparse


# Audiveris 安装路径
AUDIVERIS_PATH = r"C:\Program Files\Audiveris\Audiveris.exe"


def find_audiveris() -> str:
    """查找 Audiveris 可执行文件路径。"""
    if os.path.exists(AUDIVERIS_PATH):
        return AUDIVERIS_PATH

    # 尝试常见安装位置
    possible_paths = [
        r"C:\Program Files\Audiveris\Audiveris.exe",
        r"C:\Program Files (x86)\Audiveris\Audiveris.exe",
        os.path.expanduser(r"~\AppData\Local\Programs\Audiveris\Audiveris.exe"),
    ]

    for path in possible_paths:
        if os.path.exists(path):
            return path

    raise FileNotFoundError(
        "Audiveris not found. Please install from:\n"
        "  winget install --id Audiveris.Audiveris\n"
        "  or download from: https://github.com/Audiveris/audiveris/releases"
    )


def extract_musicxml_from_mxl(mxl_path: str, output_dir: str) -> str:
    """从 .mxl 文件中提取 MusicXML 内容。"""
    with zipfile.ZipFile(mxl_path, 'r') as z:
        for name in z.namelist():
            if name.endswith('.xml'):
                xml_path = os.path.join(output_dir, os.path.basename(name))
                with z.open(name) as src:
                    content = src.read()
                with open(xml_path, 'wb') as dst:
                    dst.write(content)
                return xml_path
    return None


def process_with_audiveris(input_path: str, output_dir: str) -> dict:
    """
    使用 Audiveris 处理输入文件。

    Returns:
        dict with keys:
            - mxl_path: path to .mxl output
            - xml_path: path to extracted .xml (if available)
            - omr_path: path to .omr project file
    """
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"Input file not found: {input_path}")

    os.makedirs(output_dir, exist_ok=True)

    audiveris_exe = find_audiveris()
    print(f"[Audiveris] Found at: {audiveris_exe}")

    # 构建命令：-batch -export -output <dir> <input>
    cmd = [
        audiveris_exe,
        '-batch',
        '-export',
        '-output', output_dir,
        input_path
    ]

    print(f"[Audiveris] Processing: {input_path}")
    print(f"[Audiveris] Command: {' '.join(cmd)}")

    # 运行 Audiveris
    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
        encoding='utf-8',
        errors='replace'
    )

    if result.returncode != 0:
        print(f"[Audiveris] STDERR:\n{result.stderr}")
        raise RuntimeError(f"Audiveris failed with code {result.returncode}")

    print(f"[Audiveris] Processing complete!")
    print(f"[Audiveris] Output directory: {output_dir}")

    # 查找生成的输出文件
    base_name = os.path.splitext(os.path.basename(input_path))[0]

    # 可能的输出文件
    mxl_path = os.path.join(output_dir, f"{base_name}.mxl")
    omr_path = os.path.join(output_dir, f"{base_name}.omr")

    # 列出输出目录内容
    print(f"[Audiveris] Files in output directory:")
    for f in os.listdir(output_dir):
        print(f"  - {f}")

    return {
        'mxl_path': mxl_path if os.path.exists(mxl_path) else None,
        'xml_path': None,  # 稍后提取
        'omr_path': omr_path if os.path.exists(omr_path) else None,
        'output_dir': output_dir,
        'base_name': base_name
    }


def main():
    parser = argparse.ArgumentParser(
        description="OMR using Audiveris - Convert sheet music images to MusicXML",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python audiveris_to_musescore.py "D:\\scores\\sheet.png" "D:\\output"
    python audiveris_to_musescore.py "D:\\scores\\piece.pdf"
        """
    )
    parser.add_argument('input', help="Input image (PNG/JPG) or PDF file")
    parser.add_argument('output_dir', nargs='?', default=None,
                        help="Output directory (default: same as input)")

    args = parser.parse_args()

    input_path = os.path.abspath(args.input)

    if args.output_dir:
        output_dir = os.path.abspath(args.output_dir)
    else:
        # 默认输出到输入文件同目录
        output_dir = os.path.join(os.path.dirname(input_path), 'audiveris_output')

    try:
        result = process_with_audiveris(input_path, output_dir)

        print()
        print("=" * 50)
        print("Processing complete!")
        print("=" * 50)

        if result['mxl_path']:
            print(f"MusicXML (.mxl): {result['mxl_path']}")
            # 提取 XML 以便查看
            xml_path = extract_musicxml_from_mxl(
                result['mxl_path'],
                output_dir
            )
            if xml_path:
                print(f"MusicXML (.xml):  {xml_path}")

        if result['omr_path']:
            print(f"Project (.omr):    {result['omr_path']}")

        print()
        print("Open .mxl file in MuseScore to verify and edit.")

    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()