src/web_ui.py

"""
Gradio Web UI for Ada to Python Converter with Langfuse instrumentation
"""

import gradio as gr
import tempfile
import time
import threading
import os
from pathlib import Path
from typing import Optional, Tuple, Generator, List, Dict, Any
from loguru import logger

from .agents.model_factory import create_model
from .agents.orchestrator_agent import OrchestratorAgent
from .utils.zip_utils import ZipHandler
from .utils.project_converter import AdaProjectConverter
from .logger.config import configure_logging

# Langfuse instrumentation using OpenTelemetry
try:
    from langfuse import get_client
    LANGFUSE_AVAILABLE = True
except ImportError:
    LANGFUSE_AVAILABLE = False


class AdaConverterUI:
    """Gradio UI for Ada to Python conversion."""
    
    def __init__(self):
        configure_logging(log_level="INFO", log_to_file=False)
        self.model = None
        self.zip_handler = ZipHandler()
        self.current_conversion_results = None
        self.current_extracted_path = None
        
        # Initialize Langfuse client if available
        self.langfuse_client = None
        if LANGFUSE_AVAILABLE:
            self._setup_langfuse()
        
        # Initialize model
        try:
            self.model = create_model()
            self.orchestrator = OrchestratorAgent(self.model)
            logger.info("Model initialized successfully for UI")
        except Exception as e:
            logger.error("Failed to initialize model for UI", error=str(e))
            raise
    
    def _setup_langfuse(self):
        """Setup Langfuse client for the UI using OpenTelemetry"""
        try:
            langfuse_secret_key = os.getenv("LANGFUSE_SECRET_KEY")
            langfuse_public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
            langfuse_host = os.getenv("LANGFUSE_HOST", "https://cloud.langfuse.com")
            
            if langfuse_secret_key and langfuse_public_key:
                # Set environment variables for automatic client initialization
                os.environ["LANGFUSE_SECRET_KEY"] = langfuse_secret_key
                os.environ["LANGFUSE_PUBLIC_KEY"] = langfuse_public_key
                os.environ["LANGFUSE_HOST"] = langfuse_host
                
                # Initialize the client using the new OpenTelemetry-based SDK
                self.langfuse_client = get_client()
                logger.info("Langfuse client initialized for Web UI with OpenTelemetry")
            else:
                logger.info("Langfuse not configured for Web UI")
        except Exception as e:
            logger.warning("Failed to setup Langfuse client", error=str(e))
    
    def convert_single_ada_file(self, file_path: str) -> Tuple[str, str, str]:
        """
        Convert a single Ada file to Python using the orchestrator.
        
        Args:
            file_path: Path to the Ada file to convert
            
        Returns:
            Tuple of (analysis_markdown, python_code, unit_test_code)
        """
        # Start Langfuse span if available
        if LANGFUSE_AVAILABLE and self.langfuse_client:
            with self.langfuse_client.start_as_current_span(name="convert_single_ada_file") as span:
                return self._convert_single_ada_file_impl(file_path, span)
        else:
            return self._convert_single_ada_file_impl(file_path, None)
    
    def _convert_single_ada_file_impl(self, file_path: str, span) -> Tuple[str, str, str]:
        """Implementation of single Ada file conversion with optional Langfuse tracing."""
        try:
            file_path = Path(file_path)
            
            # Add Langfuse metadata
            if span:
                span.update(
                    input=str(file_path),
                    metadata={
                        "file_path": str(file_path),
                        "file_extension": file_path.suffix.lower(),
                        "interface": "web_ui",
                        "operation": "single_file_conversion"
                    }
                )
            
            # Check if it's an Ada file
            if not file_path.suffix.lower() in ['.ads', '.adb', '.ada']:
                error_msg = "❌ Selected file is not an Ada source file"
                if span:
                    span.update(
                        output=error_msg,
                        level="ERROR"
                    )
                return error_msg, "", ""
            
            # Check if file exists
            if not file_path.exists():
                error_msg = "❌ Selected file does not exist"
                if span:
                    span.update(
                        output=error_msg,
                        level="ERROR"
                    )
                return error_msg, "", ""
            
            # Read the Ada code
            try:
                ada_code = file_path.read_text(encoding='utf-8')
            except UnicodeDecodeError:
                try:
                    ada_code = file_path.read_text(encoding='latin-1')
                except Exception as e:
                    error_msg = f"❌ Could not read file: {str(e)}"
                    if span:
                        span.update(
                            output=error_msg,
                            level="ERROR"
                        )
                    return error_msg, "", ""
            
            logger.info("Converting single Ada file", file_path=str(file_path), code_length=len(ada_code))
            
            # Update Langfuse with Ada code metadata
            if span:
                span.update(
                    metadata={
                        "file_path": str(file_path),
                        "file_extension": file_path.suffix.lower(),
                        "ada_code_length": len(ada_code),
                        "ada_lines": ada_code.count('\n') + 1,
                        "interface": "web_ui",
                        "operation": "single_file_conversion"
                    }
                )
            
            # Use orchestrator to convert
            analysis, python_code, unit_tests = self.orchestrator.convert_ada_to_python(ada_code)
            
            logger.success("Single Ada file converted successfully", file_path=str(file_path))
            
            # Update Langfuse with results
            if span:
                span.update(
                    output={
                        "analysis_length": len(analysis),
                        "python_code_length": len(python_code),
                        "unit_tests_length": len(unit_tests),
                        "success": True
                    }
                )
            
            return analysis, python_code, unit_tests
            
        except Exception as e:
            error_msg = f"❌ **Conversion Failed**\n\nError: {str(e)}"
            logger.exception("Failed to convert single Ada file", file_path=file_path, error=str(e))
            
            # Update Langfuse with error
            if span:
                span.update(
                    output=error_msg,
                    level="ERROR"
                )
            
            return error_msg, "", ""
    
    
    def validate_zip_file(self, file_path: Optional[str]) -> Tuple[bool, str]:
        """
        Validate that the uploaded file is a valid ZIP containing Ada files.
        
        Args:
            file_path: Path to the uploaded file
            
        Returns:
            Tuple of (is_valid, message)
        """
        if not file_path:
            self.current_extracted_path = None
            return False, "No file uploaded"
        
        file_path = Path(file_path)
        
        # Check if it's a ZIP file
        if not self.zip_handler.is_zip_file(str(file_path)):
            self.current_extracted_path = None
            return False, f"❌ Invalid file: Must be a ZIP file (got: {file_path.suffix})"
        
        try:
            # Try to extract and check for Ada files
            ada_project_path, session_dir = self.zip_handler.extract_ada_project(str(file_path))
            ada_files = list(Path(ada_project_path).rglob("*.ad[s|b]")) + list(Path(ada_project_path).rglob("*.ada"))
            
            if not ada_files:
                self.current_extracted_path = None
                return False, "❌ No Ada files found in ZIP archive"
            
            # Store the extracted path for file explorer
            self.current_extracted_path = ada_project_path
            return True, f"✅ Valid Ada project ZIP with {len(ada_files)} Ada files"
            
        except Exception as e:
            self.current_extracted_path = None
            return False, f"❌ Error validating ZIP: {str(e)}"
    
    def convert_ada_project(self, file_path: str, progress: gr.Progress) -> Generator[Tuple[str, str, bool, Optional[str]], None, None]:
        """
        Convert Ada project from ZIP file with progress updates.
        
        Args:
            file_path: Path to the uploaded ZIP file
            progress: Gradio progress tracker
            
        Yields:
            Tuples of (status_message, progress_text, download_enabled, download_file)
        """
        # Start Langfuse span if available
        if LANGFUSE_AVAILABLE and self.langfuse_client:
            with self.langfuse_client.start_as_current_span(name="convert_ada_project") as span:
                span.update(
                    input=file_path,
                    metadata={
                        "file_path": file_path,
                        "interface": "web_ui",
                        "operation": "project_conversion"
                    }
                )
                yield from self._convert_ada_project_impl(file_path, progress, span)
        else:
            yield from self._convert_ada_project_impl(file_path, progress, None)
    
    def _convert_ada_project_impl(self, file_path: str, progress: gr.Progress, span) -> Generator[Tuple[str, str, bool, Optional[str]], None, None]:
        """Implementation of Ada project conversion with optional Langfuse tracing."""
        try:
            # Validate file first
            is_valid, validation_msg = self.validate_zip_file(file_path)
            if not is_valid:
                yield validation_msg, "❌ Validation failed", False, None
                return
            
            yield validation_msg, "🔍 Validation complete", False, None
            time.sleep(0.5)  # Brief pause for UI feedback
            
            # Extract the Ada project
            progress(0.1, desc="Extracting Ada project...")
            yield "🗂️ Extracting Ada project from ZIP...", "📦 Extracting...", False, None
            
            ada_project_path, session_dir = self.zip_handler.extract_ada_project(str(file_path))
            
            progress(0.2, desc="Setting up conversion...")
            yield "⚙️ Setting up conversion process...", "⚙️ Initializing...", False, None
            
            # Create project converter
            project_converter = AdaProjectConverter(self.model)
            
            # Start conversion with progress tracking
            output_base = f"{Path(file_path).stem}_python"
            generator = project_converter.convert_project(ada_project_path, output_base)
            
            # Track conversion progress
            try:
                while True:
                    conv_progress, message = next(generator)
                    
                    # Map conversion progress (0-100) to UI progress (0.2-0.9)
                    ui_progress = 0.2 + (conv_progress / 100.0) * 0.7
                    progress(ui_progress, desc=message)
                    
                    yield f"🔄 {message}", f"[{conv_progress:5.1f}%] {message}", False, None
                    
            except StopIteration as e:
                # Conversion completed, get results
                result = e.value
                self.current_conversion_results = result
            
            progress(0.95, desc="Creating Python project ZIP...")
            yield "📦 Creating Python project ZIP file...", "📦 Packaging...", False, None
            
            # Create ZIP file for the Python project
            python_zip_path = self.zip_handler.create_python_project_zip(result['python_project'])
            
            progress(1.0, desc="Conversion completed!")
            
            # Final success message
            success_msg = f"""
🎉 **Conversion Completed Successfully!**

📊 **Results:**
- ✅ Successfully converted: {result['successful_conversions']}/{result['total_files']} files
- 📂 Python project: `{Path(result['python_project']).name}`
- 📦 ZIP file ready for download

🚀 **Next Steps:**
1. Download the Python project ZIP file below
2. Extract and explore the converted Python code
3. Run tests with: `pytest tests/ -v`
"""
            
            yield success_msg, "✅ Complete!", True, python_zip_path
            
        except Exception as e:
            logger.exception("Conversion failed in UI", error=str(e))
            error_msg = f"❌ **Conversion Failed**\n\nError: {str(e)}\n\nPlease check your ZIP file and try again."
            
            # Update Langfuse with error
            if span:
                span.update(
                    output=error_msg,
                    level="ERROR"
                )
            
            yield error_msg, "❌ Failed", False, None
    
    def handle_file_upload(self, file_path: str) -> Tuple[str, str, bool, Optional[str]]:
        """
        Handle file upload and validation.
        
        Args:
            file_path: Path to uploaded file
            
        Returns:
            Tuple of (status_message, progress_text, start_button_enabled, extracted_path)
        """
        if not file_path:
            return "Please upload a ZIP file containing an Ada project", "", False, None
        
        is_valid, message = self.validate_zip_file(file_path)
        
        if is_valid and self.current_extracted_path:
            return message, "Ready to convert", True, self.current_extracted_path
        else:
            return message, "Invalid file", False, None
    
    def create_interface(self) -> gr.Interface:
        """Create the Gradio interface."""
        
        with gr.Blocks(
            title="Ada to Python Converter",
            theme=gr.themes.Soft(),
            css="""
            .gradio-container {
                max-width: 900px !important;
            }
            .upload-container {
                border: 2px dashed #ccc;
                border-radius: 10px;
                padding: 20px;
                text-align: center;
                background-color: #f9f9f9;
            }
            """
        ) as interface:
            
            gr.Markdown("""
            # 🔄 Ada to Python Converter
            
            Upload a ZIP file containing your Ada project and convert it to Python with comprehensive unit tests.
            
            ## How to use:
            1. **Upload** a ZIP file containing `.ads`, `.adb`, or `.ada` files
            2. **Start** the conversion process
            3. **Download** the converted Python project ZIP when ready
            """)
            
            with gr.Row():
                with gr.Column(scale=2):
                    # File upload with ZIP restriction
                    file_upload = gr.File(
                        label="📁 Upload Ada Project ZIP File",
                        file_types=[".zip"],
                        file_count="single",
                        elem_classes=["upload-container"]
                    )
                    
                    # Status message
                    status_display = gr.Markdown("Please upload a ZIP file to begin")
                    
                    # Progress indicator
                    progress_text = gr.Textbox(
                        label="Progress",
                        value="",
                        interactive=False,
                        visible=True
                    )
                
                with gr.Column(scale=1):
                    # Start conversion button
                    start_button = gr.Button(
                        "🚀 Start Conversion",
                        variant="primary",
                        interactive=False,
                        size="lg"
                    )
                    
                    # Download button
                    download_button = gr.File(
                        label="📥 Download Python Project",
                        visible=False,
                        interactive=False
                    )
            
            # File explorer section
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 📂 Extracted Project Structure")
                    file_explorer = gr.FileExplorer(
                        root_dir="/",  # Will be updated dynamically
                        label="Browse Ada Project Files",
                        visible=False,
                        interactive=True,
                        file_count="single",
                        height=400
                    )
                    
                    # Selected file display
                    selected_file_display = gr.Textbox(
                        label="Selected File",
                        value="No file selected",
                        interactive=False,
                        visible=False
                    )
            
            # Single file conversion results (tabbed interface)
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### 🔄 Single File Conversion Results")
                    with gr.Tabs(visible=False) as conversion_tabs:
                        with gr.Tab("📋 Analysis"):
                            analysis_output = gr.Markdown("", label="Analysis Report")
                        with gr.Tab("🐍 Python Code"):
                            python_output = gr.Code("", language="python", label="Generated Python Code")
                        with gr.Tab("🧪 Unit Tests"):
                            tests_output = gr.Code("", language="python", label="Generated Unit Tests")
            
            # Conversion output area (for project conversion)
            with gr.Row():
                conversion_output = gr.Markdown(
                    "",
                    label="Project Conversion Output",
                    visible=True
                )
            
            # Event handlers
            def update_ui_on_upload(file_path):
                status_msg, progress_msg, button_enabled, extracted_path = self.handle_file_upload(file_path)
                
                # Show file explorer if we have valid extracted path
                show_explorer = extracted_path is not None
                
                return [
                    status_msg,
                    progress_msg,
                    gr.update(interactive=button_enabled),
                    gr.update(visible=False),  # Hide download button
                    gr.update(visible=show_explorer, root_dir=extracted_path if extracted_path else "/"),  # Update file explorer
                    gr.update(visible=show_explorer, value="No file selected" if show_explorer else "")  # Update selected file display
                ]
            
            def run_conversion(file_path, progress=gr.Progress()):
                if not file_path:
                    return "Please upload a file first", "", gr.update(interactive=False), gr.update(visible=False)
                
                # Generator for progress updates
                for status_msg, progress_msg, download_enabled, download_file in self.convert_ada_project(file_path, progress):
                    yield [
                        status_msg,
                        progress_msg,
                        gr.update(interactive=False),  # Keep button disabled during conversion
                        gr.update(visible=download_enabled, value=download_file if download_enabled else None)
                    ]
            
            def on_file_select(selected_file):
                """Handle file selection in the explorer"""
                if not selected_file:
                    return [
                        "No file selected",
                        gr.update(visible=False),  # Hide conversion tabs
                        "",  # Clear analysis
                        "",  # Clear Python code
                        ""   # Clear unit tests
                    ]
                
                file_path = Path(selected_file)
                display_msg = f"Selected: {file_path.name}"
                
                # Check if it's an Ada file
                if file_path.suffix.lower() in ['.ads', '.adb', '.ada']:
                    display_msg += " (Converting...)"
                    
                    # Convert the Ada file
                    analysis, python_code, unit_tests = self.convert_single_ada_file(selected_file)
                    
                    display_msg = f"Selected: {file_path.name} ✅ Converted"
                    
                    return [
                        display_msg,
                        gr.update(visible=True),   # Show conversion tabs
                        analysis,                  # Update analysis
                        python_code,              # Update Python code
                        unit_tests                # Update unit tests
                    ]
                else:
                    return [
                        display_msg + " (Not an Ada file)",
                        gr.update(visible=False),  # Hide conversion tabs
                        "",  # Clear analysis
                        "",  # Clear Python code
                        ""   # Clear unit tests
                    ]
            
            # Wire up events
            file_upload.change(
                fn=update_ui_on_upload,
                inputs=[file_upload],
                outputs=[status_display, progress_text, start_button, download_button, file_explorer, selected_file_display]
            )
            
            # Wire up file explorer selection
            file_explorer.change(
                fn=on_file_select,
                inputs=[file_explorer],
                outputs=[selected_file_display, conversion_tabs, analysis_output, python_output, tests_output]
            )
            
            start_button.click(
                fn=run_conversion,
                inputs=[file_upload],
                outputs=[conversion_output, progress_text, start_button, download_button]
            )
        
        return interface


def launch_ui(share: bool = False, server_port: int = 7860):
    """Launch the Gradio UI."""
    try:
        ui = AdaConverterUI()
        interface = ui.create_interface()
        
        print("🚀 Starting Ada to Python Converter Web UI...")
        print(f"📡 Server will be available at: http://localhost:{server_port}")
        
        if share:
            print("🌐 Public URL will be generated (share=True)")
        
        interface.launch(
            share=share,
            server_port=server_port,
            server_name="0.0.0.0",
            show_api=False,
            favicon_path=None
        )
        
    except Exception as e:
        logger.exception("Failed to launch UI", error=str(e))
        print(f"❌ Failed to launch UI: {e}")


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Launch Ada to Python Converter Web UI")
    parser.add_argument("--share", action="store_true", help="Create public URL (default: False)")
    parser.add_argument("--port", type=int, default=7860, help="Server port (default: 7860)")
    
    args = parser.parse_args()
    
    launch_ui(share=args.share, server_port=args.port)