import gradio as gr
import tempfile
from pathlib import Path
import json
import torch
import time
import re
import threading
from queue import Queue
import zipfile
from datetime import datetime

from detection_ocr import DetectionOCR
from repair_head import RepairHead
from semantic_evaluation import SemanticEvaluationPipeline

# ============================================================
# Device Configuration
# ============================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[System] Using device: {DEVICE}")

# ============================================================
# Lazy Singletons
# ============================================================
_ocr = None
_repair = None
_semantic = None

def get_ocr():
    global _ocr
    if _ocr is None:
        print("[INIT] Loading DetectionOCR")
        _ocr = DetectionOCR(device=DEVICE)
    return _ocr

def get_repair():
    global _repair
    if _repair is None:
        print("[INIT] Loading RepairHead (CPU)")
        _repair = RepairHead(device="cpu")
    return _repair

def get_semantic():
    global _semantic
    if _semantic is None:
        print("[INIT] Loading Semantic Evaluation Pipeline")
        _semantic = SemanticEvaluationPipeline(
            base_model="Qwen/Qwen2.5-Math-1.5B-Instruct",
            lora_dir=None,
            device="cpu",
        )
    return _semantic

# ============================================================
# Helpers
# ============================================================
KNOWN_ENV = ["equation", "align", "gather", "multline", "array", "cases"]

def wrap_equation(expr: str) -> str:
    expr = expr.strip()
    if any(expr.startswith(f"\\begin{{{e}}}") for e in KNOWN_ENV):
        return expr
    if expr.startswith("$$") and expr.endswith("$$"):
        return expr
    if expr.startswith("\\[") and expr.endswith("\\]"):
        return expr
    return f"\\begin{{equation}}\n{expr}\n\\end{{equation}}"

def is_reasonable(expr: str) -> bool:
    if not expr or not isinstance(expr, str):
        return False
    if expr.lower() in ["noise", "error"]:
        return False
    if len(expr) < 3 or len(expr) > 2000:
        return False
    return True

def save_progress_log(outdir: Path, stage: str, content: str):
    """Save progress logs to files"""
    log_file = outdir / f"{stage}_log.txt"
    with open(log_file, "a", encoding="utf-8") as f:
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        f.write(f"[{timestamp}] {content}\n")

def create_visualization_pdf(outdir: Path, repaired_data: dict, semantic_results: list):
    """创建可视化PDF，按页面顺序显示公式"""
    try:
        from fpdf import FPDF
        
        pdf = FPDF()
        pdf.set_auto_page_break(auto=True, margin=15)
        
        # 添加标题页
        pdf.add_page()
        pdf.set_font('Arial', 'B', 16)
        pdf.cell(0, 10, 'PDF Formula Processing Results', 0, 1, 'C')
        pdf.ln(10)
        
        pdf.set_font('Arial', '', 12)
        pdf.cell(0, 10, f'Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}', 0, 1, 'C')
        pdf.ln(20)
        
        # 统计信息
        total_pages = len(repaired_data)
        total_formulas = sum(len(page_data) for page_data in repaired_data.values())
        valid_formulas = sum(len([f for f in page_data if f['repaired'] not in ['noise', 'error']]) 
                           for page_data in repaired_data.values())
        
        pdf.set_font('Arial', 'B', 14)
        pdf.cell(0, 10, 'Processing Statistics', 0, 1, 'L')
        pdf.ln(5)
        
        pdf.set_font('Arial', '', 12)
        pdf.cell(0, 8, f'• Total Pages: {total_pages}', 0, 1, 'L')
        pdf.cell(0, 8, f'• Total Formulas Detected: {total_formulas}', 0, 1, 'L')
        pdf.cell(0, 8, f'• Valid Formulas Repaired: {valid_formulas}', 0, 1, 'L')
        pdf.cell(0, 8, f'• Semantic Evaluations: {len(semantic_results)}', 0, 1, 'L')
        pdf.ln(10)
        
        # 按页面处理
        page_keys = sorted(repaired_data.keys(), key=lambda x: int(re.search(r'page_(\d+)', x).group(1)) if re.search(r'page_(\d+)', x) else 0)
        
        for page_key in page_keys:
            formulas = repaired_data[page_key]
            valid_formulas = [f for f in formulas if f['repaired'] not in ['noise', 'error']]
            
            if not valid_formulas:
                continue
            
            pdf.add_page()
            pdf.set_font('Arial', 'B', 14)
            pdf.cell(0, 10, f'{page_key.upper()}', 0, 1, 'L')
            pdf.set_font('Arial', '', 10)
            pdf.cell(0, 6, f'Valid Formulas: {len(valid_formulas)} / {len(formulas)}', 0, 1, 'L')
            pdf.ln(5)
            
            for idx, formula_data in enumerate(valid_formulas):
                formula = formula_data['repaired']
                original = formula_data['raw_latex'][:150] + "..." if len(formula_data['raw_latex']) > 150 else formula_data['raw_latex']
                
                # 公式标题
                pdf.set_font('Arial', 'B', 10)
                pdf.cell(0, 6, f'Formula {idx+1}:', 0, 1, 'L')
                
                # 原始公式
                pdf.set_font('Arial', 'I', 8)
                pdf.cell(20, 4, 'Original:', 0, 0, 'L')
                pdf.set_font('Courier', '', 8)
                pdf.multi_cell(0, 4, f' {original}', 0, 'L')
                
                # 修复后公式
                pdf.set_font('Arial', 'I', 8)
                pdf.cell(20, 4, 'Repaired:', 0, 0, 'L')
                pdf.set_font('Courier', '', 8)
                formula_display = formula[:150] + "..." if len(formula) > 150 else formula
                pdf.multi_cell(0, 4, f' {formula_display}', 0, 'L')
                
                # 查找对应的语义结果
                sem_result = None
                for sem in semantic_results:
                    if sem.get('original_formula') == formula_data['raw_latex'] or sem.get('fixed_formula') == formula:
                        sem_result = sem
                        break
                
                if sem_result:
                    verdict = sem_result.get('semantic_verdict', 'unknown')
                    confidence = sem_result.get('confidence', 0)
                    pdf.set_font('Arial', '', 9)
                    verdict_color = {
                        'equivalent': (0, 100, 0),  # 绿色
                        'contradiction': (220, 0, 0),  # 红色
                        'non_verifiable': (255, 140, 0),  # 橙色
                        'garbled': (128, 128, 128),  # 灰色
                    }.get(verdict, (0, 0, 0))  # 黑色默认
                    
                    pdf.set_text_color(*verdict_color)
                    pdf.cell(0, 5, f'Verdict: {verdict} (Confidence: {confidence:.2f})', 0, 1, 'L')
                    pdf.set_text_color(0, 0, 0)  # 重置为黑色
                    
                    if sem_result.get('explanation'):
                        pdf.set_font('Arial', 'I', 8)
                        explanation = sem_result['explanation'][:200] + "..." if len(sem_result['explanation']) > 200 else sem_result['explanation']
                        pdf.multi_cell(0, 4, f'Explanation: {explanation}', 0, 'L')
                
                pdf.ln(3)
                
                # 每5个公式后添加分隔线
                if (idx + 1) % 5 == 0 and idx + 1 < len(valid_formulas):
                    pdf.line(10, pdf.get_y(), 200, pdf.get_y())
                    pdf.ln(5)
        
        # 保存PDF
        pdf_path = outdir / "visualization_report.pdf"
        pdf.output(str(pdf_path))
        
        return pdf_path
        
    except Exception as e:
        print(f"Error creating visualization PDF: {e}")
        # 如果fpdf不可用，创建简单的文本报告
        try:
            report_path = outdir / "visualization_report.txt"
            with open(report_path, 'w', encoding='utf-8') as f:
                f.write("=" * 60 + "\n")
                f.write("PDF FORMULA PROCESSING RESULTS\n")
                f.write("=" * 60 + "\n")
                f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
                
                # 统计信息
                total_pages = len(repaired_data)
                total_formulas = sum(len(page_data) for page_data in repaired_data.values())
                valid_formulas = sum(len([f for f in page_data if f['repaired'] not in ['noise', 'error']]) 
                                   for page_data in repaired_data.values())
                
                f.write("PROCESSING STATISTICS:\n")
                f.write(f"• Total Pages: {total_pages}\n")
                f.write(f"• Total Formulas Detected: {total_formulas}\n")
                f.write(f"• Valid Formulas Repaired: {valid_formulas}\n")
                f.write(f"• Semantic Evaluations: {len(semantic_results)}\n\n")
                
                f.write("=" * 60 + "\n\n")
                
                # 按页面显示公式
                page_keys = sorted(repaired_data.keys(), key=lambda x: int(re.search(r'page_(\d+)', x).group(1)) if re.search(r'page_(\d+)', x) else 0)
                
                for page_key in page_keys:
                    formulas = repaired_data[page_key]
                    valid_formulas_list = [f for f in formulas if f['repaired'] not in ['noise', 'error']]
                    
                    if not valid_formulas_list:
                        continue
                    
                    f.write(f"\n{page_key.upper()}\n")
                    f.write("-" * 40 + "\n")
                    f.write(f"Valid Formulas: {len(valid_formulas_list)} / {len(formulas)}\n\n")
                    
                    for idx, formula_data in enumerate(valid_formulas_list):
                        formula = formula_data['repaired']
                        original = formula_data['raw_latex']
                        
                        f.write(f"Formula {idx+1}:\n")
                        f.write(f"  Original: {original}\n")
                        f.write(f"  Repaired: {formula}\n")
                        
                        # 查找语义结果
                        sem_result = None
                        for sem in semantic_results:
                            if sem.get('original_formula') == original or sem.get('fixed_formula') == formula:
                                sem_result = sem
                                break
                        
                        if sem_result:
                            verdict = sem_result.get('semantic_verdict', 'unknown')
                            confidence = sem_result.get('confidence', 0)
                            f.write(f"  Verdict: {verdict} (Confidence: {confidence:.2f})\n")
                            
                            if sem_result.get('explanation'):
                                explanation = sem_result['explanation']
                                f.write(f"  Explanation: {explanation}\n")
                        
                        f.write("\n")
            
            return report_path
        except Exception as e2:
            print(f"Fallback text report also failed: {e2}")
            return None

# ============================================================
# Stream Manager for Real-time Updates
# ============================================================
class StreamManager:
    """Manage real-time streaming to multiple outputs"""
    def __init__(self):
        self.ocr_queue = Queue()
        self.repair_queue = Queue()
        self.semantic_queue = Queue()
        self.summary_queue = Queue()
        
    def ocr_stream(self, message: str):
        self.ocr_queue.put(message)
        if hasattr(self, 'current_outdir'):
            save_progress_log(self.current_outdir, "ocr", message)
        
    def repair_stream(self, message: str):
        self.repair_queue.put(message)
        if hasattr(self, 'current_outdir'):
            save_progress_log(self.current_outdir, "repair", message)
        
    def semantic_stream(self, message: str):
        self.semantic_queue.put(message)
        if hasattr(self, 'current_outdir'):
            save_progress_log(self.current_outdir, "semantic", message)
        
    def summary_stream(self, message: str):
        self.summary_queue.put(message)
        if hasattr(self, 'current_outdir'):
            save_progress_log(self.current_outdir, "summary", message)
        
    def set_output_dir(self, outdir: Path):
        self.current_outdir = outdir

# ============================================================
# Main Processing Pipeline
# ============================================================
def process_pipeline(pdf_file, stream_mgr: StreamManager):
    """Main processing pipeline with streaming to different outputs"""
    
    if pdf_file is None:
        stream_mgr.summary_stream("❌ No PDF uploaded.")
        return None, None, None, None, None, []
    
    try:
        # Create output directory
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        tmpdir = Path(tempfile.mkdtemp(prefix=f"pdf_pipeline_{timestamp}_"))
        outdir = tmpdir / "output"
        outdir.mkdir(exist_ok=True)
        
        stream_mgr.set_output_dir(outdir)
        
        stream_mgr.summary_stream(f"📁 Output directory: {outdir}")
        stream_mgr.summary_stream("=" * 60)
        
        # ============ STAGE 1: OCR ============
        stream_mgr.ocr_stream("🚀 STAGE 1: OCR Detection Started")
        stream_mgr.ocr_stream("Loading OCR model...")
        
        ocr = get_ocr()
        stream_mgr.ocr_stream("Processing PDF pages...")
        
        ocr.process_pdf(pdf_file, output_dir=outdir)
        
        stream_mgr.ocr_stream("✅ OCR Processing Completed")
        stream_mgr.summary_stream("✓ Stage 1: OCR completed")
        
        # ============ STAGE 2: Merge Results ============
        stream_mgr.ocr_stream("\n🔄 Merging page results...")
        
        merged = {}
        page_files = sorted(outdir.glob("page_*_results.json"))
        
        if not page_files:
            stream_mgr.ocr_stream("❌ No page results found!")
            stream_mgr.summary_stream("❌ Error: No page results found")
            return None, None, None, None, None, []
        
        for p in page_files:
            m = re.search(r"(page_\d+)_results\.json", p.name)
            if not m:
                continue
            page = m.group(1)
            try:
                data = json.loads(p.read_text("utf-8"))
                formulas = data.get("formulas", [])
                merged[page] = {
                    "formulas": formulas,
                    "total": len(formulas)
                }
                stream_mgr.ocr_stream(f"  {page}: {len(formulas)} formulas")
            except Exception as e:
                stream_mgr.ocr_stream(f"  ⚠️ Error loading {page}: {str(e)}")
                merged[page] = {"formulas": [], "total": 0}
        
        # Save overall OCR results
        overall_json = outdir / "overall_ocr_results.json"
        overall_json.write_text(json.dumps(merged, indent=2), "utf-8")
        
        stream_mgr.ocr_stream(f"✅ Saved: {overall_json.name}")
        stream_mgr.summary_stream(f"✓ Stage 2: Merged {len(merged)} pages")
        
        # ============ STAGE 3: Repair Formulas ============
        stream_mgr.repair_stream("🔧 STAGE 3: LaTeX Repair Started")
        stream_mgr.repair_stream("Loading Repair model...")
        
        repair = get_repair()
        
        repaired_all = {}
        repaired_eqs = []
        total_formulas = sum(v["total"] for v in merged.values())
        
        stream_mgr.repair_stream(f"Total formulas to repair: {total_formulas}")
        stream_mgr.repair_stream("=" * 40)
        
        processed = 0
        for page, content in merged.items():
            stream_mgr.repair_stream(f"\n📄 Processing {page}:")
            repaired_all[page] = []
            
            for i, fm in enumerate(content["formulas"]):
                processed += 1
                raw = fm.get("raw_latex", "")
                
                # Show progress every 5 formulas or first/last
                if processed <= 5 or processed >= total_formulas - 5 or processed % 10 == 0:
                    raw_preview = raw[:100] + "..." if len(raw) > 100 else raw
                    stream_mgr.repair_stream(f"  [{processed}/{total_formulas}] {raw_preview}")
                
                try:
                    repaired = repair.repair(raw)
                    
                    if is_reasonable(repaired):
                        repaired_all[page].append({
                            "raw_latex": raw,
                            "repaired": repaired
                        })
                        repaired_eqs.append(repaired)
                        if processed <= 3:  # Show first few repairs
                            stream_mgr.repair_stream(f"    → {repaired[:80]}...")
                    else:
                        repaired_all[page].append({
                            "raw_latex": raw,
                            "repaired": "noise"
                        })
                except Exception as e:
                    repaired_all[page].append({
                        "raw_latex": raw,
                        "repaired": f"error: {str(e)}"
                    })
                    stream_mgr.repair_stream(f"    ⚠️ Error: {str(e)}")
        
        # Save repaired results
        repaired_json = outdir / "overall_ocr_repaired.json"
        repaired_json.write_text(json.dumps(repaired_all, indent=2), "utf-8")
        
        stream_mgr.repair_stream(f"\n✅ Repair completed: {len(repaired_eqs)} valid formulas")
        stream_mgr.summary_stream(f"✓ Stage 3: Repaired {len(repaired_eqs)}/{total_formulas} formulas")
        
        # ============ STAGE 4: Semantic Evaluation ============
        stream_mgr.semantic_stream("🧠 STAGE 4: Semantic Evaluation Started")
        stream_mgr.semantic_stream("Loading Semantic model...")
        
        semantic = get_semantic()
        
        semantic_results = []
        stream_mgr.semantic_stream(f"Evaluating {len(repaired_eqs)} formulas...")
        stream_mgr.semantic_stream("=" * 40)
        
        for i, eq in enumerate(repaired_eqs):
            try:
                # Show progress
                if i < 3 or i >= len(repaired_eqs) - 3 or i % 5 == 0:
                    eq_preview = eq[:80] + "..." if len(eq) > 80 else eq
                    stream_mgr.semantic_stream(f"[{i+1}/{len(repaired_eqs)}] {eq_preview}")
                
                res = semantic.evaluate_formula(eq)
                res['original_formula'] = eq  # 添加原始公式
                semantic_results.append(res)
                
                # Show verdict for first few
                if i < 5:
                    stream_mgr.semantic_stream(f"    Verdict: {res.get('semantic_verdict', 'unknown')}")
                    stream_mgr.semantic_stream(f"    Confidence: {res.get('confidence', 0):.2f}")
                
            except Exception as e:
                stream_mgr.semantic_stream(f"⚠️ Error evaluating formula {i+1}: {str(e)}")
                semantic_results.append({
                    'original_formula': eq,
                    'error': str(e),
                    'semantic_verdict': 'evaluation_error',
                    'confidence': 0.0
                })
        
        # Save semantic results
        semantic_json = outdir / "semantic_results.json"
        semantic_json.write_text(json.dumps(semantic_results, indent=2), "utf-8")
        
        stream_mgr.semantic_stream(f"\n✅ Semantic evaluation completed: {len(semantic_results)} formulas")
        stream_mgr.summary_stream(f"✓ Stage 4: Semantically evaluated {len(semantic_results)} formulas")
        
        # ============ STAGE 5: Generate Visualization PDF ============
        stream_mgr.summary_stream("\n📊 Generating visualization PDF...")
        
        # Generate LaTeX file
        merge_tex = outdir / "merge.tex"
        with merge_tex.open("w", encoding="utf-8") as f:
            f.write("\\documentclass{article}\n")
            f.write("\\usepackage{amsmath, amssymb, bm}\n")
            f.write("\\begin{document}\n\n")
            f.write("% Auto-generated from PDF formula extraction pipeline\n")
            f.write("% Generated on: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n\n")
            for eq in repaired_eqs:
                f.write(wrap_equation(eq) + "\n\n")
            f.write("\\end{document}\n")
        
        stream_mgr.summary_stream("✓ Generated LaTeX file")
        
        # Create visualization PDF
        visualization_pdf = create_visualization_pdf(outdir, repaired_all, semantic_results)
        
        if visualization_pdf and visualization_pdf.exists():
            stream_mgr.summary_stream(f"✓ Created visualization report: {visualization_pdf.name}")
        else:
            stream_mgr.summary_stream("⚠️ Could not create visualization PDF")
            visualization_pdf = None
        
        # ============ STAGE 6: Generate Statistics ============
        stats = {
            "total_pages": len(merged),
            "total_formulas": total_formulas,
            "repaired_formulas": len(repaired_eqs),
            "semantic_evaluated": len(semantic_results),
            "output_directory": str(outdir),
            "processing_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
        
        stats_json = outdir / "processing_stats.json"
        stats_json.write_text(json.dumps(stats, indent=2), "utf-8")
        
        # Final summary
        stream_mgr.summary_stream("=" * 60)
        stream_mgr.summary_stream("🎉 PROCESSING COMPLETED SUCCESSFULLY!")
        stream_mgr.summary_stream("=" * 60)
        stream_mgr.summary_stream(f"📊 SUMMARY STATISTICS:")
        stream_mgr.summary_stream(f"  • Pages processed: {len(merged)}")
        stream_mgr.summary_stream(f"  • Formulas detected: {total_formulas}")
        stream_mgr.summary_stream(f"  • Formulas repaired: {len(repaired_eqs)}")
        stream_mgr.summary_stream(f"  • Formulas evaluated: {len(semantic_results)}")
        stream_mgr.summary_stream(f"  • Output directory: {outdir}")
        stream_mgr.summary_stream("=" * 60)
        stream_mgr.summary_stream("📁 Generated files:")
        stream_mgr.summary_stream(f"  1. {overall_json.name} - OCR results")
        stream_mgr.summary_stream(f"  2. {repaired_json.name} - Repaired formulas")
        stream_mgr.summary_stream(f"  3. {semantic_json.name} - Semantic evaluation")
        stream_mgr.summary_stream(f"  4. {merge_tex.name} - Combined LaTeX")
        if visualization_pdf:
            stream_mgr.summary_stream(f"  5. {visualization_pdf.name} - Visualization report")
        stream_mgr.summary_stream(f"  6. {stats_json.name} - Processing statistics")
        stream_mgr.summary_stream("=" * 60)
        
        return overall_json, repaired_json, semantic_json, merge_tex, visualization_pdf, semantic_results
        
    except Exception as e:
        import traceback
        error_msg = f"❌ CRITICAL ERROR:\n{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        stream_mgr.summary_stream(error_msg)
        return None, None, None, None, None, []

# ============================================================
# Simplified Streaming Handler
# ============================================================
def process_with_streaming(pdf_file):
    """Simplified streaming handler without threading issues"""
    if not pdf_file:
        yield "", "", "", "❌ No PDF uploaded", "Error", None, None, None, None, None, []
        return
    
    # Initialize stream manager
    stream_mgr = StreamManager()
    
    # Initialize result variables
    result_container = {
        'overall_json': None,
        'repaired_json': None,
        'semantic_json': None,
        'merge_tex': None,
        'visualization_pdf': None,
        'sem_results': []
    }
    
    # Process function
    def run_processing():
        """Run the pipeline and store results"""
        results = process_pipeline(pdf_file, stream_mgr)
        if results:
            result_container['overall_json'] = results[0]
            result_container['repaired_json'] = results[1]
            result_container['semantic_json'] = results[2]
            result_container['merge_tex'] = results[3]
            result_container['visualization_pdf'] = results[4]
            result_container['sem_results'] = results[5]
    
    # Start processing in a thread
    process_thread = threading.Thread(target=run_processing)
    process_thread.start()
    
    # Initialize log contents
    ocr_log_content = ""
    repair_log_content = ""
    semantic_log_content = ""
    summary_log_content = ""
    
    # Stream updates while processing
    last_update = time.time()
    while process_thread.is_alive():
        # Collect new messages
        while not stream_mgr.ocr_queue.empty():
            ocr_log_content += stream_mgr.ocr_queue.get() + "\n"
        
        while not stream_mgr.repair_queue.empty():
            repair_log_content += stream_mgr.repair_queue.get() + "\n"
        
        while not stream_mgr.semantic_queue.empty():
            semantic_log_content += stream_mgr.semantic_queue.get() + "\n"
        
        while not stream_mgr.summary_queue.empty():
            summary_log_content += stream_mgr.summary_queue.get() + "\n"
        
        # Update UI periodically
        current_time = time.time()
        if current_time - last_update > 0.1:
            yield (
                ocr_log_content[-5000:],  # Limit size
                repair_log_content[-5000:],
                semantic_log_content[-5000:],
                summary_log_content[-5000:],
                "Processing...",
                None, None, None, None, None, []
            )
            last_update = current_time
        
        time.sleep(0.05)
    
    # Wait for thread to complete
    process_thread.join()
    
    # Collect any remaining messages
    while not stream_mgr.ocr_queue.empty():
        ocr_log_content += stream_mgr.ocr_queue.get() + "\n"
    
    while not stream_mgr.repair_queue.empty():
        repair_log_content += stream_mgr.repair_queue.get() + "\n"
    
    while not stream_mgr.semantic_queue.empty():
        semantic_log_content += stream_mgr.semantic_queue.get() + "\n"
    
    while not stream_mgr.summary_queue.empty():
        summary_log_content += stream_mgr.summary_queue.get() + "\n"
    
    # Prepare download files
    ocr_file = str(result_container['overall_json']) if result_container['overall_json'] else None
    repair_file = str(result_container['repaired_json']) if result_container['repaired_json'] else None
    semantic_file = str(result_container['semantic_json']) if result_container['semantic_json'] else None
    tex_file = str(result_container['merge_tex']) if result_container['merge_tex'] else None
    pdf_file_dl = str(result_container['visualization_pdf']) if result_container['visualization_pdf'] else None
    
    # Final update
    yield (
        ocr_log_content[-10000:],
        repair_log_content[-10000:],
        semantic_log_content[-10000:],
        summary_log_content[-10000:],
        "Completed",
        ocr_file,
        repair_file,
        semantic_file,
        tex_file,
        pdf_file_dl,
        result_container['sem_results']
    )

# ============================================================
# Gradio UI with Multi-Stream Display
# ============================================================
with gr.Blocks(title="PDF Formula Pipeline", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 📘 PDF Formula Processing Pipeline
    **Real-time streaming display with visualization report**
    """)
    
    # Input Section
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Upload PDF")
            pdf_input = gr.File(
                label="Upload PDF File",
                file_types=[".pdf"],
                height=100
            )
            
            with gr.Row():
                run_btn = gr.Button("Start Processing", variant="primary", size="lg")
                clear_btn = gr.Button("Clear All", variant="secondary")
            
            gr.Markdown("### Processing Status")
            status_text = gr.Textbox(
                label="Current Status",
                value="Ready",
                interactive=False
            )
    
    # Multi-Stream Display Section
    gr.Markdown("### 🎯 Real-time Processing Logs")
    
    with gr.Tabs():
        with gr.TabItem("📄 OCR Process"):
            ocr_log = gr.Textbox(
                label="OCR Stage Log",
                lines=15,
                interactive=False,
                autoscroll=True
            )
        
        with gr.TabItem("🔧 Repair Process"):
            repair_log = gr.Textbox(
                label="Repair Stage Log",
                lines=15,
                interactive=False,
                autoscroll=True
            )
        
        with gr.TabItem("🧠 Semantic Process"):
            semantic_log = gr.Textbox(
                label="Semantic Stage Log",
                lines=15,
                interactive=False,
                autoscroll=True
            )
        
        with gr.TabItem("📊 Summary"):
            summary_log = gr.Textbox(
                label="Processing Summary",
                lines=20,
                interactive=False,
                autoscroll=True
            )
    
    # Results Download Section
    gr.Markdown("### 📥 Download Results")
    
    with gr.Row():
        with gr.Column(scale=1):
            ocr_json_dl = gr.File(label="OCR Results (JSON)", visible=True)
            repaired_json_dl = gr.File(label="Repaired Formulas (JSON)", visible=True)
        
        with gr.Column(scale=1):
            semantic_json_dl = gr.File(label="Semantic Results (JSON)", visible=True)
            merge_tex_dl = gr.File(label="Combined LaTeX (.tex)", visible=True)
    
    with gr.Row():
        visualization_pdf_dl = gr.File(label="Visualization Report", visible=True)
    
    # Formula Details Section
    gr.Markdown("### 🔍 Formula Details")
    
    with gr.Row():
        formula_idx = gr.Number(
            label="Formula Index (0 to preview)",
            minimum=0,
            maximum=1000,
            step=1,
            precision=0,
            value=0
        )
        preview_btn = gr.Button("Show Details", variant="secondary")
    
    formula_details = gr.JSON(
        label="Formula Details",
        visible=True
    )
    
    # Hidden state for semantic results
    semantic_state = gr.State([])
    
    # Preview function
    def show_formula_details(idx, sem_results):
        if not sem_results or idx < 0 or idx >= len(sem_results):
            return {}
        
        result = sem_results[int(idx)]
        
        # 格式化输出
        formatted_result = {
            "Formula": result.get('original_formula', 'N/A'),
            "Semantic Verdict": result.get('semantic_verdict', 'unknown'),
            "Confidence": f"{result.get('confidence', 0):.2f}" if isinstance(result.get('confidence'), (int, float)) else "N/A",
            "Source": result.get('source', 'N/A'),
            "Syntactically Correct": result.get('syntactically_correct', 'N/A'),
            "Mathematically Correct": result.get('mathematically_correct', 'N/A'),
        }
        
        if result.get('explanation'):
            formatted_result["Explanation"] = result['explanation']
        
        if result.get('suggested_fix'):
            formatted_result["Suggested Fix"] = result['suggested_fix']
        
        if result.get('error'):
            formatted_result["Error"] = result['error']
        
        return formatted_result
    
    # Clear function
    def clear_all():
        return [
            None,  # pdf_input
            "",    # ocr_log
            "",    # repair_log
            "",    # semantic_log
            "",    # summary_log
            "Ready",  # status_text
            None,  # ocr_json_dl
            None,  # repaired_json_dl
            None,  # semantic_json_dl
            None,  # merge_tex_dl
            None,  # visualization_pdf_dl
            [],    # semantic_state
            0,     # formula_idx
            {}     # formula_details
        ]
    
    # Connect events
    outputs = [
        ocr_log,
        repair_log,
        semantic_log,
        summary_log,
        status_text,
        ocr_json_dl,
        repaired_json_dl,
        semantic_json_dl,
        merge_tex_dl,
        visualization_pdf_dl,
        semantic_state
    ]
    
    run_btn.click(
        fn=process_with_streaming,
        inputs=[pdf_input],
        outputs=outputs
    )
    
    preview_btn.click(
        fn=show_formula_details,
        inputs=[formula_idx, semantic_state],
        outputs=[formula_details]
    )
    
    formula_idx.change(
        fn=show_formula_details,
        inputs=[formula_idx, semantic_state],
        outputs=[formula_details]
    )
    
    clear_btn.click(
        fn=clear_all,
        outputs=outputs + [formula_idx, formula_details]
    )

# ============================================================
# Main Execution
# ============================================================
if __name__ == "__main__":
    print("=" * 60)
    print("PDF Formula Processing Pipeline - Visualization Version")
    print(f"Device: {DEVICE}")
    print("=" * 60)
    
    # 尝试导入fpdf，如果失败则使用文本报告
    try:
        from fpdf import FPDF
        print("✓ fpdf is available for PDF generation")
    except ImportError:
        print("⚠️ fpdf not available, will use text report instead")
        print("   To install: pip install fpdf")
    
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=False,
        show_error=True
    )