import gradio as gr
import subprocess
import threading
import time

def get_md_content(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        return f"Error: {file_path} not found."
    except Exception as e:
        return f"An error occurred: {e}"

def run_script():
    """Function to run the fine-tuning script and stream output."""
    process = subprocess.Popen(
        ['python3', 'fine_tune_improved.py'],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        universal_newlines=True
    )
    output = ""
    for line in process.stdout:
        output += line
        yield output
    process.wait()

# JavaScript to find and render Mermaid diagrams
js_script = """
() => {
    function initMermaidAndConvert() {
        // Wait for mermaid to be available
        if (typeof mermaid === 'undefined') {
            console.log('Mermaid not loaded yet, retrying...');
            setTimeout(initMermaidAndConvert, 100);
            return;
        }
        
        console.log('Mermaid loaded successfully');
        // Initialize mermaid
        mermaid.initialize({ 
            startOnLoad: false, 
            theme: 'default',
            securityLevel: 'loose'
        });
        
        function convertMermaidCodeBlocks() {
            console.log('Converting Mermaid code blocks...');
            let processedCount = 0;
            
            // Look for pre blocks that contain mermaid syntax
            document.querySelectorAll('pre').forEach((pre, index) => {
                const code = pre.querySelector('code');
                if (code && !pre.classList.contains('mermaid-processed')) {
                    const text = code.textContent.trim();
                    // Check if it contains mermaid syntax
                    const isMermaid = text.includes('graph ') || 
                                    text.includes('flowchart ') || 
                                    text.includes('subgraph ') ||
                                    text.startsWith('graph') ||
                                    text.startsWith('flowchart') ||
                                    text.includes('classDiagram') ||
                                    text.includes('sequenceDiagram');
                    
                    if (isMermaid) {
                        console.log(`Found Mermaid diagram ${processedCount + 1}:`, text.substring(0, 50) + '...');
                        pre.classList.add('mermaid');
                        pre.classList.add('mermaid-processed');
                        pre.textContent = text;
                        processedCount++;
                    }
                }
            });
            
            console.log(`Processed ${processedCount} Mermaid diagrams`);
            // Run Mermaid
            try {
                mermaid.run();
            } catch (e) {
                console.log('Mermaid rendering error:', e);
            }
        }

        // Use a MutationObserver to re-run the conversion when Gradio updates the page
        const observer = new MutationObserver((mutations) => {
            // A simple debounce to avoid excessive re-renders
            clearTimeout(window.mermaidTimeout);
            window.mermaidTimeout = setTimeout(convertMermaidCodeBlocks, 100);
        });
        observer.observe(document.body, { childList: true, subtree: true });

        // Initial run
        convertMermaidCodeBlocks();
    }
    
    // Start the initialization
    initMermaidAndConvert();
}
"""

# HTML to include the Mermaid.js library
head_script = '<script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>'

with gr.Blocks(theme=gr.themes.Soft(), head=head_script, js=js_script) as demo:
    gr.Markdown("# 微调技术分享")

    with gr.Tabs():
        with gr.TabItem("分享大纲"):
            gr.Markdown(get_md_content("outline.md"))
        
        with gr.TabItem("核心技术概览"):
            gr.Markdown(get_md_content("presentation.md"))

        with gr.TabItem("LoRA & QLoRA 深度解析"):
            gr.Markdown(get_md_content("lora_qlora_deep_dive.md"))



        with gr.TabItem("动手实战：模型微调"):
            with gr.Row():
                start_button = gr.Button("开始微调", variant="primary")
            
            log_output = gr.Textbox(
                label="训练日志",
                interactive=False,
                lines=20,
                show_copy_button=True
            )
            start_button.click(fn=run_script, outputs=log_output)

if __name__ == "__main__":
    demo.launch()