import asyncio
import re
import html as html_module
import markdown
try:
    import bleach
    _HAS_BLEACH = True
except ImportError:
    _HAS_BLEACH = False
from nonebot import logger

async def markdown_to_image(markdown_text: str, output_path: str, browser=None):
    """将 Markdown 转换为 HTML 并使用 Puppeteer 截图。"""
    page = None
    should_close_browser = False
    try:
        # Convert markdown to HTML. The markdown library handles special chars safely.
        # Note: do NOT html.escape() before markdown.markdown() - it breaks markdown syntax.
        html_content = markdown.markdown(markdown_text, extensions=["fenced_code", "tables"])
        # Sanitize to prevent XSS from malicious AI responses
        allowed_tags = [
            'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p', 'br', 'hr',
            'ul', 'ol', 'li', 'blockquote', 'pre', 'code', 'span',
            'table', 'thead', 'tbody', 'tr', 'th', 'td',
            'strong', 'em', 'b', 'i', 'u', 'a', 'img', 'div',
        ]
        allowed_attrs = {'a': ['href', 'title'], 'img': ['src', 'alt', 'title']}
        if _HAS_BLEACH:
            html_content = bleach.clean(html_content, tags=allowed_tags, attributes=allowed_attrs)
        else:
            logger.warning("[chatai] bleach 未安装，跳过 HTML 净化（请运行 pip install bleach）")
        
        # 使用传入的浏览器实例或创建新的
        if browser is None:
            from pyppeteer import launch
            browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])
            should_close_browser = True
            
        page = await browser.newPage()
        page.setDefaultNavigationTimeout(15000)
        
        # 设置页面样式，使内容更美观
        await page.setContent(f"""
        <html>
            <head>
                <style>
                    body {{
                        font-family: "PingFang SC", "Microsoft YaHei", sans-serif;
                        line-height: 1.6;
                        padding: 30px;
                        max-width: 800px;
                        margin: 0 auto;
                        background-color: transparent;
                        color: #333;
                    }}
                    .container {{
                        background-color: #ffffff;
                        border-radius: 15px;
                        box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
                        padding: 25px;
                        overflow: hidden;
                    }}
                    p {{
                        margin-bottom: 16px;
                    }}
                    code {{
                        font-family: "SFMono-Regular", Consolas, "Liberation Mono", Menlo, monospace;
                        background-color: #f5f7f9;
                        padding: 2px 6px;
                        border-radius: 4px;
                        font-size: 0.9em;
                    }}
                    pre {{
                        background-color: #f5f7f9;
                        padding: 15px;
                        border-radius: 8px;
                        overflow-x: auto;
                        margin: 20px 0;
                    }}
                    pre code {{
                        background-color: transparent;
                        padding: 0;
                    }}
                    h1, h2, h3, h4, h5, h6 {{
                        margin-top: 24px;
                        margin-bottom: 16px;
                        font-weight: 600;
                        line-height: 1.25;
                    }}
                    h1 {{
                        font-size: 1.8em;
                        border-bottom: 1px solid #eaecef;
                        padding-bottom: 0.3em;
                    }}
                    h2 {{
                        font-size: 1.5em;
                        border-bottom: 1px solid #eaecef;
                        padding-bottom: 0.3em;
                    }}
                    blockquote {{
                        padding: 0 1em;
                        color: #6a737d;
                        border-left: 0.25em solid #dfe2e5;
                        margin: 16px 0;
                    }}
                    ul, ol {{
                        padding-left: 2em;
                        margin-bottom: 16px;
                    }}
                    img {{
                        max-width: 100%;
                        border-radius: 8px;
                    }}
                    a {{
                        color: #0366d6;
                        text-decoration: none;
                    }}
                    a:hover {{
                        text-decoration: underline;
                    }}
                    table {{
                        border-collapse: collapse;
                        width: 100%;
                        margin: 16px 0;
                    }}
                    table th, table td {{
                        padding: 8px 12px;
                        border: 1px solid #dfe2e5;
                    }}
                    table th {{
                        background-color: #f6f8fa;
                    }}
                </style>
            </head>
            <body>
                <div class="container">
                    {html_content}
                </div>
            </body>
        </html>
        """)
        
        # 等待内容渲染完成
        try:
            await asyncio.wait_for(page.waitForNavigation({'waitUntil': 'networkidle0'}), timeout=10)
        except Exception:
            pass  # rendering may already be complete
        
        # 获取内容尺寸并设置视口
        dimensions = await page.evaluate('''() => {
            const container = document.querySelector('.container');
            return {
                width: container.offsetWidth + 60, // 加上 body 的 padding
                height: container.offsetHeight + 60
            }
        }''')
        
        # 设置视口大小
        await page.setViewport({
            'width': dimensions['width'],
            'height': dimensions['height'],
            'deviceScaleFactor': 2  # 提高图片清晰度
        })
        
        # 截图，使用透明背景
        await page.screenshot({
            'path': output_path, 
            'omitBackground': True,  # 透明背景
            'clip': {
                'x': 0,
                'y': 0,
                'width': dimensions['width'],
                'height': dimensions['height']
            }
        })
        
        # 关闭页面
        await page.close()
        
        # 如果是我们创建的浏览器实例，则关闭它
        if should_close_browser:
            await browser.close()
            
    except Exception as e:
        # 确保资源被释放
        if page is not None:
            try:
                await page.close()
            except Exception:
                pass
        if should_close_browser and browser is not None:
            try:
                await browser.close()
            except Exception:
                pass
        raise  # 重新抛出异常以便上层处理