import sys
import asyncio
from pathlib import Path
from playwright.async_api import async_playwright
import subprocess

HTML_TMPL = """
<!DOCTYPE html>
<html>
<head>
  <meta charset="utf-8"/>
  <script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
  <style>
    body { margin: 0; padding: 20px; }
    #container { width: 1800px; }
    .mermaid { font-size: 16px; }
  </style>
  <script>
    mermaid.initialize({ startOnLoad: true, theme: 'default' });
  </script>
</head>
<body>
  <div id="container">
    <div class="mermaid" id="diagram">DIAGRAM_PLACEHOLDER</div>
  </div>
  </body>
</html>
"""

async def render(mermaid_text: str, out_png: Path):
    html = HTML_TMPL.replace("DIAGRAM_PLACEHOLDER", mermaid_text)
    tmp = out_png.with_suffix('.html')
    tmp.write_text(html, encoding='utf-8')
    async def _do_render():
        async with async_playwright() as p:
            browser = await p.chromium.launch()
            page = await browser.new_page(viewport={"width": 2000, "height": 1200})
            await page.goto(tmp.resolve().as_uri(), wait_until="domcontentloaded", timeout=120000)
            await page.wait_for_selector('#diagram', state='attached', timeout=120000)
            await page.evaluate("""
            try { mermaid.initialize({ startOnLoad: false, theme: 'default' }); } catch (e) {}
            """
            )
            await page.evaluate("""
            (async () => {
              try {
                await mermaid.run({ querySelector: '.mermaid' });
                const svg = document.querySelector('.mermaid svg');
                if (svg) { svg.style.width = '1800px'; svg.style.height = 'auto'; svg.style.maxWidth = 'none'; }
              } catch (e) {}
            })();
            """
            )
            await page.wait_for_selector('.mermaid svg', state='attached', timeout=120000)
            elem = await page.query_selector('#container')
            await elem.screenshot(path=str(out_png))
            await browser.close()
    try:
        await _do_render()
    except Exception:
        try:
            subprocess.run([sys.executable, '-m', 'playwright', 'install', 'chromium'], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
            await _do_render()
        finally:
            tmp.unlink(missing_ok=True)
    else:
        tmp.unlink(missing_ok=True)

class _BatchRenderer:
    def __init__(self, pw, browser, page):
        self._pw = pw
        self._browser = browser
        self._page = page
    @classmethod
    async def create(cls):
        pw = await async_playwright().start()
        browser = await pw.chromium.launch()
        page = await browser.new_page(viewport={"width": 2000, "height": 1200})
        return cls(pw, browser, page)
    async def render(self, mermaid_text: str, out_png: Path):
        html = HTML_TMPL.replace("DIAGRAM_PLACEHOLDER", mermaid_text)
        tmp = out_png.with_suffix('.html')
        tmp.write_text(html, encoding='utf-8')
        await self._page.goto(tmp.resolve().as_uri(), wait_until="domcontentloaded", timeout=120000)
        await self._page.wait_for_selector('#diagram', state='attached', timeout=120000)
        await self._page.evaluate("""
        try { mermaid.initialize({ startOnLoad: false, theme: 'default' }); } catch (e) {}
        """
        )
        await self._page.evaluate("""
        (async () => {
          try {
            await mermaid.run({ querySelector: '.mermaid' });
            const svg = document.querySelector('.mermaid svg');
            if (svg) { svg.style.width = '1800px'; svg.style.height = 'auto'; svg.style.maxWidth = 'none'; }
          } catch (e) {}
        })();
        """
        )
        await self._page.wait_for_selector('.mermaid svg', state='attached', timeout=120000)
        elem = await self._page.query_selector('#container')
        await elem.screenshot(path=str(out_png))
        tmp.unlink(missing_ok=True)
    async def close(self):
        try:
            await self._browser.close()
        finally:
            try:
                await self._pw.stop()
            except Exception:
                pass

_batch = None

def render_fast(mermaid_text: str, out_png: Path):
    asyncio.run(render(mermaid_text, out_png))

def close_fast():
    try:
        pass
    except Exception:
        pass

def main():
    import argparse
    parser = argparse.ArgumentParser(
        prog='HOS_M2F.2png',
        description='Mermaid 转 PNG',
        epilog=(
            '示例:\n'
            '  HOS_M2F.2png input.mmd output.png\n'
            '相关命令:\n'
            '  HOS_M2F.2docx input.md output.docx [-p paper|patent|softcopy] [-nm] [-nc] [-q]\n'
            '  HOS_M2F.2html input.md output.html [-p wx|csdn]\n'
            '  HOS_M2F.2pdf  input.md output.pdf\n'
            '  HOS_M2F.2json input.md output.json\n'
            '  HOS_M2F.2xml  input.md output.xml\n'
        ),
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument('input', help='输入 Mermaid 文件路径')
    parser.add_argument('output', help='输出 PNG 文件路径')
    args = parser.parse_args()
    inp = Path(args.input)
    outp = Path(args.output)
    text = inp.read_text(encoding='utf-8')
    asyncio.run(render(text, outp))
    print(f'Rendered: {inp} -> {outp}')

if __name__ == '__main__':
    main()
