app.py

import gradio as gr
import sys
import os
import tempfile
import shutil
import ast
import time
import subprocess
import re
from typing import List, Dict, Optional, Tuple, Any
from py2puml.py2puml import py2puml
from plantuml import PlantUML
import pyan
from pathlib import Path
from utils import setup_testing_space, verify_testing_space, cleanup_testing_space

if os.name == "nt":  # nt == Windows
    graphviz_bin = r"C:\\Program Files\\Graphviz\\bin"
    if graphviz_bin not in os.environ["PATH"]:
        os.environ["PATH"] += os.pathsep + graphviz_bin


def generate_call_graph_with_pyan3(
    python_code: str, filename: str = "analysis"
) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]:
    """Generate call graph using pyan3 and return DOT content, PNG path, and structured data.

    Args:
        python_code: The Python code to analyze
        filename: Base filename for temporary files

    Returns:
        Tuple of (dot_content, png_path, structured_data)
    """
    if not python_code.strip():
        return None, None, {}

    # Create unique filename using timestamp
    timestamp = str(int(time.time() * 1000))
    unique_filename = f"{filename}_{timestamp}"

    # Paths
    testing_dir = os.path.join(os.getcwd(), "inputs")
    code_file = os.path.join(testing_dir, f"{unique_filename}.py")

    try:
        # Write Python code to file
        with open(code_file, "w", encoding="utf-8") as f:
            f.write(python_code)

        print(f"📊 Generating call graph for: {unique_filename}.py")

        try:

            dot_content = pyan.create_callgraph(
                filenames=[str(code_file)],
                format="dot",
                colored=True,
                grouped=True,
                annotated=True,
            )

            png_path = None
            with tempfile.TemporaryDirectory() as temp_dir:
                dot_file = os.path.join(temp_dir, f"{unique_filename}.dot")
                temp_png = os.path.join(temp_dir, f"{unique_filename}.png")

                # Write DOT content to file
                with open(dot_file, "w", encoding="utf-8") as f:
                    f.write(dot_content)

                # Generate PNG using dot command
                dot_cmd = ["dot", "-Tpng", dot_file, "-o", temp_png]

                try:
                    subprocess.run(dot_cmd, check=True, timeout=30)

                    if os.path.exists(temp_png):
                        # Copy to permanent location
                        permanent_dir = os.path.join(os.getcwd(), "temp_diagrams")
                        os.makedirs(permanent_dir, exist_ok=True)
                        png_path = os.path.join(
                            permanent_dir, f"callgraph_{unique_filename}.png"
                        )
                        shutil.copy2(temp_png, png_path)
                        print(f"🎨 Call graph PNG saved: {os.path.basename(png_path)}")

                except subprocess.SubprocessError as e:
                    print(f"⚠️ Graphviz PNG generation failed: {e}")
                    # Continue without PNG, DOT content is still useful

            # Parse DOT content for structured data
            structured_data = parse_call_graph_data(dot_content)

            return dot_content, png_path, structured_data

        except subprocess.TimeoutExpired:
            print("⚠️ pyan3 analysis timed out, trying simplified approach...")
            return try_fallback_analysis(python_code, unique_filename)
        except subprocess.SubprocessError as e:
            print(f"⚠️ pyan3 execution failed: {e}, trying fallback...")
            return try_fallback_analysis(python_code, unique_filename)

    except Exception as e:
        print(f"❌ Call graph generation error: {e}")
        return None, None, {"error": str(e)}

    finally:
        # Clean up temporary file
        if os.path.exists(code_file):
            try:
                os.remove(code_file)
                print(f"🧹 Cleaned up analysis file: {unique_filename}.py")
            except Exception as e:
                print(f"⚠️ Could not remove analysis file: {e}")


def parse_call_graph_data(dot_content: str) -> Dict[str, Any]:
    """Parse pyan3 DOT output into structured function call data.

    Args:
        dot_content: DOT format string from pyan3

    Returns:
        Dictionary with parsed call graph information
    """
    if not dot_content:
        return {}

    try:
        # Extract nodes (functions/classes)
        node_pattern = r'"([^"]+)"\s*\['
        nodes = re.findall(node_pattern, dot_content)

        # Extract edges (function calls)
        edge_pattern = r'"([^"]+)"\s*->\s*"([^"]+)"'
        edges = re.findall(edge_pattern, dot_content)

        # Build function call mapping
        call_graph = {}
        called_by = {}

        for caller, callee in edges:
            if caller not in call_graph:
                call_graph[caller] = []
            call_graph[caller].append(callee)

            if callee not in called_by:
                called_by[callee] = []
            called_by[callee].append(caller)

        # Calculate metrics
        function_metrics = {}
        for node in nodes:
            out_degree = len(call_graph.get(node, []))
            in_degree = len(called_by.get(node, []))

            function_metrics[node] = {
                "calls_made": out_degree,
                "called_by_count": in_degree,
                "calls_to": call_graph.get(node, []),
                "called_by": called_by.get(node, []),
            }

        return {
            "nodes": nodes,
            "edges": edges,
            "total_functions": len(nodes),
            "total_calls": len(edges),
            "call_graph": call_graph,
            "function_metrics": function_metrics,
        }

    except Exception as e:
        return {"parse_error": str(e)}


def try_fallback_analysis(
    python_code: str, unique_filename: str
) -> Tuple[Optional[str], Optional[str], Dict[str, Any]]:
    """Fallback analysis when pyan3 fails - basic function call detection.

    Args:
        python_code: The Python code to analyze
        unique_filename: Unique filename for this analysis

    Returns:
        Tuple of (None, None, fallback_analysis_data)
    """
    print("🔄 Using fallback analysis approach...")

    try:
        import ast
        import re

        tree = ast.parse(python_code)
        functions = []
        calls = []

        # Extract function definitions
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                functions.append(node.name)

        # Simple regex-based call detection (fallback approach)
        for func in functions:
            # Look for calls to this function
            pattern = rf"\b{re.escape(func)}\s*\("
            if re.search(pattern, python_code):
                calls.append(("unknown", func))

        return (
            None,
            None,
            {
                "fallback": True,
                "functions_detected": functions,
                "total_functions": len(functions),
                "total_calls": len(calls),
                "info": f"Fallback analysis: detected {len(functions)} functions",
                "function_metrics": {
                    func: {
                        "calls_made": 0,
                        "called_by_count": 0,
                        "calls_to": [],
                        "called_by": [],
                    }
                    for func in functions
                },
            },
        )

    except Exception as e:
        return None, None, {"error": f"Fallback analysis also failed: {str(e)}"}


def analyze_function_complexity(python_code: str) -> Dict[str, Any]:
    """Analyze function complexity using AST.

    Args:
        python_code: The Python code to analyze

    Returns:
        Dictionary with function complexity metrics
    """
    if not python_code.strip():
        return {}

    try:
        tree = ast.parse(python_code)
        function_analysis = {}

        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef):
                # Calculate cyclomatic complexity (simplified)
                complexity = 1  # Base complexity

                for child in ast.walk(node):
                    if isinstance(
                        child,
                        (
                            ast.If,
                            ast.While,
                            ast.For,
                            ast.Try,
                            ast.ExceptHandler,
                            ast.With,
                            ast.Assert,
                        ),
                    ):
                        complexity += 1
                    elif isinstance(child, ast.BoolOp):
                        complexity += len(child.values) - 1

                # Count lines of code
                lines = (
                    node.end_lineno - node.lineno + 1
                    if hasattr(node, "end_lineno")
                    else 0
                )

                # Extract parameters
                params = [arg.arg for arg in node.args.args]

                # Check for docstring
                has_docstring = (
                    len(node.body) > 0
                    and isinstance(node.body[0], ast.Expr)
                    and isinstance(node.body[0].value, ast.Constant)
                    and isinstance(node.body[0].value.value, str)
                )

                function_analysis[node.name] = {
                    "complexity": complexity,
                    "lines_of_code": lines,
                    "parameter_count": len(params),
                    "parameters": params,
                    "has_docstring": has_docstring,
                    "line_start": node.lineno,
                    "line_end": getattr(node, "end_lineno", node.lineno),
                }

        return function_analysis

    except Exception as e:
        return {"error": str(e)}


def generate_diagram(python_code: str, filename: str = "diagram") -> Optional[str]:
    """Generate a UML class diagram from Python code.

    Args:
        python_code: The Python code to analyze and convert to UML
        filename: Optional name for the generated diagram file

    Returns:
        Path to the generated PNG diagram image or None if failed
    """
    if not python_code.strip():
        return None

    print(f"🔄 Processing code for diagram generation...")

    # Clean testing space (ensure only __init__.py exists)
    cleanup_testing_space()

    # Verify clean state
    if not verify_testing_space():
        print("⚠️ testing_space verification failed, recreating...")
        setup_testing_space()
        cleanup_testing_space()

    # Create unique filename using timestamp
    timestamp = str(int(time.time() * 1000))  # millisecond timestamp
    unique_filename = f"{filename}_{timestamp}"

    # Paths
    testing_dir = os.path.join(os.getcwd(), "inputs")
    code_file = os.path.join(testing_dir, f"{unique_filename}.py")

    # Use PlantUML web service for rendering
    server = PlantUML(url="http://www.plantuml.com/plantuml/img/")

    try:
        # Write Python code to file in testing_space
        with open(code_file, "w", encoding="utf-8") as f:
            f.write(python_code)

        print(f"📝 Created temporary file: inputs/{unique_filename}.py")

        # Generate PlantUML content using py2puml (no sys.path manipulation needed)
        print(f"📝 Generating PlantUML content...")
        puml_content_lines = py2puml(
            os.path.join(
                testing_dir, unique_filename
            ),  # path to the .py file (without extension)
            f"inputs.{unique_filename}",  # module name
        )
        puml_content = "".join(puml_content_lines)

        if not puml_content.strip():
            print("⚠️ No UML content generated - check if your code contains classes")
            return None

        # Create temporary directory for PlantUML processing
        with tempfile.TemporaryDirectory() as temp_dir:
            # Save PUML file
            puml_file = os.path.join(temp_dir, f"{unique_filename}.puml")
            with open(puml_file, "w", encoding="utf-8") as f:
                f.write(puml_content)

            print(f"🎨 Rendering diagram...")
            # Generate PNG
            output_png = os.path.join(temp_dir, f"{unique_filename}.png")
            server.processes_file(puml_file, outfile=output_png)

            if os.path.exists(output_png):
                print("✅ Diagram generated successfully!")
                # Copy to a permanent location for Gradio to serve
                permanent_dir = os.path.join(os.getcwd(), "temp_diagrams")
                os.makedirs(permanent_dir, exist_ok=True)
                permanent_path = os.path.join(
                    permanent_dir, f"{filename}_{hash(python_code) % 10000}.png"
                )
                shutil.copy2(output_png, permanent_path)
                return permanent_path
            else:
                print("❌ Failed to generate PNG")
                return None

    except Exception as e:
        print(f"❌ Error: {e}")
        return None

    finally:
        # Always clean up the temporary .py file
        if os.path.exists(code_file):
            try:
                os.remove(code_file)
                print(f"🧹 Cleaned up temporary file: {unique_filename}.py")
            except Exception as e:
                print(f"⚠️ Could not remove temporary file: {e}")


def analyze_code_structure(python_code: str) -> str:
    """Return a Markdown report with complexity metrics and recommendations.

    Args:
        python_code: The Python code to analyze

    Returns:
        Comprehensive analysis report in markdown format
    """
    if not python_code.strip():
        return "No code provided for analysis."

    try:
        # Basic AST analysis
        tree = ast.parse(python_code)
        classes = []
        functions = []
        imports = []

        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                methods = []
                attributes = []

                for item in node.body:
                    if isinstance(item, ast.FunctionDef):
                        methods.append(item.name)
                    elif isinstance(item, ast.Assign):
                        for target in item.targets:
                            if isinstance(target, ast.Name):
                                attributes.append(target.id)

                # Check for inheritance
                parents = [base.id for base in node.bases if isinstance(base, ast.Name)]

                classes.append(
                    {
                        "name": node.name,
                        "methods": methods,
                        "attributes": attributes,
                        "parents": parents,
                    }
                )

            elif isinstance(node, ast.FunctionDef):
                # Check if it's a top-level function (not inside a class)
                is_method = any(
                    isinstance(parent, ast.ClassDef)
                    for parent in ast.walk(tree)
                    if hasattr(parent, "body") and node in getattr(parent, "body", [])
                )
                if not is_method:
                    functions.append(node.name)

            elif isinstance(node, (ast.Import, ast.ImportFrom)):
                if isinstance(node, ast.Import):
                    for alias in node.names:
                        imports.append(alias.name)
                else:
                    module = node.module or ""
                    for alias in node.names:
                        imports.append(
                            f"{module}.{alias.name}" if module else alias.name
                        )

        # Enhanced function complexity analysis
        function_complexity = analyze_function_complexity(python_code)

        # Call graph analysis (for files with functions)
        call_graph_data = {}
        if functions or any(classes):  # Only run if there are functions to analyze
            try:
                cleanup_testing_space()  # Ensure clean state
                dot_content, png_path, call_graph_data = generate_call_graph_with_pyan3(
                    python_code
                )
            except Exception as e:
                print(f"⚠️ Call graph analysis failed: {e}")
                call_graph_data = {"error": str(e)}

        # Build comprehensive summary
        summary = "📊 **Enhanced Code Analysis Results**\n\n"

        # === OVERVIEW SECTION ===
        summary += "## 📋 **Overview**\n"
        summary += f"• **{len(classes)}** classes found\n"
        summary += f"• **{len(functions)}** standalone functions found\n"
        summary += f"• **{len(set(imports))}** unique imports\n"

        if call_graph_data and "total_functions" in call_graph_data:
            summary += f"• **{call_graph_data['total_functions']}** total functions/methods in call graph\n"
            summary += (
                f"• **{call_graph_data['total_calls']}** function calls detected\n"
            )

        summary += "\n"

        # === CLASSES SECTION ===
        if classes:
            summary += "## 🏗️ **Classes**\n"
            for cls in classes:
                summary += f"### **{cls['name']}**\n"
                if cls["parents"]:
                    summary += f"  - **Inherits from**: {', '.join(cls['parents'])}\n"
                summary += f"  - **Methods**: {len(cls['methods'])}"
                if cls["methods"]:
                    summary += f" ({', '.join(cls['methods'])})"
                summary += "\n"
                if cls["attributes"]:
                    summary += f"  - **Attributes**: {', '.join(cls['attributes'])}\n"
                summary += "\n"

        # === STANDALONE FUNCTIONS SECTION ===
        if functions:
            summary += "## ⚙️ **Standalone Functions**\n"
            for func in functions:
                summary += f"### **{func}()**\n"

                # Add complexity metrics if available
                if func in function_complexity:
                    metrics = function_complexity[func]
                    summary += (
                        f"  - **Complexity**: {metrics['complexity']} (cyclomatic)\n"
                    )
                    summary += f"  - **Lines of Code**: {metrics['lines_of_code']}\n"
                    summary += f"  - **Parameters**: {metrics['parameter_count']}"
                    if metrics["parameters"]:
                        summary += f" ({', '.join(metrics['parameters'])})"
                    summary += "\n"
                    summary += f"  - **Has Docstring**: {'✅' if metrics['has_docstring'] else '❌'}\n"
                    summary += f"  - **Lines**: {metrics['line_start']}-{metrics['line_end']}\n"

                # Add call graph info if available
                if call_graph_data and "function_metrics" in call_graph_data:
                    if func in call_graph_data["function_metrics"]:
                        call_metrics = call_graph_data["function_metrics"][func]
                        summary += f"  - **Calls Made**: {call_metrics['calls_made']}\n"
                        if call_metrics["calls_to"]:
                            summary += (
                                f"    - Calls: {', '.join(call_metrics['calls_to'])}\n"
                            )
                        summary += f"  - **Called By**: {call_metrics['called_by_count']} functions\n"
                        if call_metrics["called_by"]:
                            summary += f"    - Called by: {', '.join(call_metrics['called_by'])}\n"

                summary += "\n"

        # === CALL GRAPH ANALYSIS ===
        if (
            call_graph_data
            and "function_metrics" in call_graph_data
            and call_graph_data["total_calls"] > 0
        ):
            summary += "## 🔗 **Function Call Analysis**\n"

            # Most called functions
            sorted_by_calls = sorted(
                call_graph_data["function_metrics"].items(),
                key=lambda x: x[1]["called_by_count"],
                reverse=True,
            )[:5]

            if sorted_by_calls and sorted_by_calls[0][1]["called_by_count"] > 0:
                summary += "**Most Called Functions:**\n"
                for func_name, metrics in sorted_by_calls:
                    if metrics["called_by_count"] > 0:
                        summary += f"• **{func_name}**: called {metrics['called_by_count']} times\n"
                summary += "\n"

            # Most complex functions (by calls made)
            sorted_by_complexity = sorted(
                call_graph_data["function_metrics"].items(),
                key=lambda x: x[1]["calls_made"],
                reverse=True,
            )[:5]

            if sorted_by_complexity and sorted_by_complexity[0][1]["calls_made"] > 0:
                summary += "**Functions Making Most Calls:**\n"
                for func_name, metrics in sorted_by_complexity:
                    if metrics["calls_made"] > 0:
                        summary += (
                            f"• **{func_name}**: makes {metrics['calls_made']} calls\n"
                        )
                summary += "\n"

        # === COMPLEXITY ANALYSIS ===
        if function_complexity:
            summary += "## 📈 **Complexity Analysis**\n"

            # Sort by complexity
            sorted_complexity = sorted(
                function_complexity.items(),
                key=lambda x: x[1]["complexity"],
                reverse=True,
            )[:5]

            summary += "**Most Complex Functions:**\n"
            for func_name, metrics in sorted_complexity:
                summary += f"• **{func_name}**: complexity {metrics['complexity']}, {metrics['lines_of_code']} lines\n"

            # Overall stats
            total_functions = len(function_complexity)
            avg_complexity = (
                sum(m["complexity"] for m in function_complexity.values())
                / total_functions
            )
            avg_lines = (
                sum(m["lines_of_code"] for m in function_complexity.values())
                / total_functions
            )
            functions_with_docs = sum(
                1 for m in function_complexity.values() if m["has_docstring"]
            )

            summary += "\n**Overall Function Metrics:**\n"
            summary += f"• **Average Complexity**: {avg_complexity:.1f}\n"
            summary += f"• **Average Lines per Function**: {avg_lines:.1f}\n"
            summary += f"• **Functions with Docstrings**: {functions_with_docs}/{total_functions} ({100*functions_with_docs/total_functions:.1f}%)\n"
            summary += "\n"

        # === IMPORTS SECTION ===
        if imports:
            summary += "## 📦 **Imports**\n"
            unique_imports = list(set(imports))
            for imp in unique_imports[:10]:  # Show first 10 imports
                summary += f"• {imp}\n"
            if len(unique_imports) > 10:
                summary += f"• ... and {len(unique_imports) - 10} more\n"
            summary += "\n"

        # === CALL GRAPH ERROR/INFO ===
        if call_graph_data and "error" in call_graph_data:
            summary += "## ⚠️ **Call Graph Analysis**\n"
            summary += f"Call graph generation failed: {call_graph_data['error']}\n\n"
        elif call_graph_data and "info" in call_graph_data:
            summary += "## 📊 **Call Graph Analysis**\n"
            summary += f"{call_graph_data['info']}\n\n"

        # === RECOMMENDATIONS ===
        summary += "## 💡 **Recommendations**\n"
        if function_complexity:
            high_complexity = [
                f for f, m in function_complexity.items() if m["complexity"] > 10
            ]
            if high_complexity:
                summary += f"• Consider refactoring high-complexity functions: {', '.join(high_complexity)}\n"

            no_docs = [
                f for f, m in function_complexity.items() if not m["has_docstring"]
            ]
            if no_docs:
                summary += f"• Add docstrings to: {', '.join(no_docs[:5])}{'...' if len(no_docs) > 5 else ''}\n"

        if call_graph_data and "function_metrics" in call_graph_data:
            isolated_functions = [
                f
                for f, m in call_graph_data["function_metrics"].items()
                if m["calls_made"] == 0 and m["called_by_count"] == 0
            ]
            if isolated_functions:
                summary += f"• Review isolated functions: {', '.join(isolated_functions[:3])}{'...' if len(isolated_functions) > 3 else ''}\n"

        return summary

    except SyntaxError as e:
        return f"❌ **Syntax Error in Python code:**\n```\n{str(e)}\n```"
    except Exception as e:
        return f"❌ **Error analyzing code:**\n```\n{str(e)}\n```"


def list_example_files() -> list:
    """List all example .py files in the examples/ directory."""
    examples_dir = os.path.join(os.getcwd(), "examples")
    if not os.path.exists(examples_dir):
        return []
    return [f for f in os.listdir(examples_dir) if f.endswith(".py")]


def get_sample_code(filename: str) -> str:
    """Return sample Python code from examples/ directory."""
    examples_dir = os.path.join(os.getcwd(), "examples")
    file_path = os.path.join(examples_dir, filename)
    with open(file_path, "r", encoding="utf-8") as f:
        return f.read()


def generate_all_diagrams(
    python_code: str, filename: str = "diagram"
) -> Tuple[Optional[str], Optional[str], str]:
    """Generate class diagram, call-graph diagram and analysis in one call.

    Args:
        python_code: The Python code to analyze
        filename: Base filename for diagrams

    Returns:
        Tuple of (uml_diagram_path, call_graph_path, analysis_text)
    """
    if not python_code.strip():
        return None, None, "No code provided for analysis."

    print("🚀 Starting comprehensive diagram generation...")

    # Step 1: Generate UML Class Diagram
    print("📊 Step 1/3: Generating UML class diagram...")
    uml_diagram_path = generate_diagram(python_code, filename)

    # Step 2: Generate Call Graph
    print("🔗 Step 2/3: Generating call graph...")
    try:
        cleanup_testing_space()
        dot_content, call_graph_path, structured_data = generate_call_graph_with_pyan3(
            python_code
        )
    except Exception as e:
        print(f"⚠️ Call graph generation failed: {e}")
        call_graph_path = None

    # Step 3: Generate Analysis
    print("📈 Step 3/3: Performing code analysis...")
    analysis_text = analyze_code_structure(python_code)

    print("✅ All diagrams and analysis completed!")

    return uml_diagram_path, call_graph_path, analysis_text


# =============================================================================
# ❶  Wrapper functions for diagram and analysis generation
# These will be connected to the UI buttons and the MCP interfaces.
# =============================================================================

def generate_class_diagram_only(python_code: str) -> Optional[str]:
    """Generates just the UML class diagram."""
    if not python_code.strip():
        gr.Warning("Input code is empty!")
        return None
    return generate_diagram(python_code)

def generate_call_graph_only(python_code: str) -> Optional[str]:
    """Generates just the call graph diagram."""
    if not python_code.strip():
        gr.Warning("Input code is empty!")
        return None
    _, png_path, _ = generate_call_graph_with_pyan3(python_code)
    return png_path

def analyze_code_only(python_code: str) -> str:
    """Generates just the code analysis report."""
    if not python_code.strip():
        gr.Warning("Input code is empty!")
        return "No code provided to analyze."
    return analyze_code_structure(python_code)
    
def generate_all_outputs(python_code: str) -> Tuple[Optional[str], Optional[str], str]:
    """Generates all three outputs: UML diagram, call graph, and analysis."""
    if not python_code.strip():
        gr.Warning("Input code is empty!")
        return None, None, "No code provided to analyze."
    
    print("🚀 Starting comprehensive generation...")
    uml_path = generate_diagram(python_code)
    _, call_graph_path, _ = generate_call_graph_with_pyan3(python_code)
    analysis_text = analyze_code_structure(python_code)
    print("✅ All outputs generated!")
    
    return uml_path, call_graph_path, analysis_text

# =============================================================================
# ❷  Four MCP-exposed Interfaces
# These are NOT rendered in the UI but are exposed as tools for agents.
# =============================================================================

iface_class = gr.Interface(
    fn=generate_class_diagram_only,
    inputs=gr.Textbox(lines=20, label="Python code"),
    outputs=gr.Image(label="UML diagram"),
    api_name="generate_class_diagram",
    description="Create a UML class diagram (PNG) from Python code.",
)

iface_call = gr.Interface(
    fn=generate_call_graph_only,
    inputs=gr.Textbox(lines=20, label="Python code"),
    outputs=gr.Image(label="Call‑graph"),
    api_name="generate_call_graph_diagram",
    description="Generate a function‑call graph (PNG) from Python code.",
)

iface_analysis = gr.Interface(
    fn=analyze_code_only,
    inputs=gr.Textbox(lines=20, label="Python code"),
    outputs=gr.Markdown(label="Analysis"),
    api_name="analyze_code_structure",
    description="Return a Markdown report with complexity metrics.",
)

iface_all = gr.Interface(
    fn=generate_all_outputs,
    inputs=gr.Textbox(lines=20, label="Python code"),
    outputs=[
        gr.Image(label="UML diagram"),
        gr.Image(label="Call‑graph"),
        gr.Markdown(label="Analysis"),
    ],
    api_name="generate_all",
    description="Run class diagram, call graph and analysis in one call.",
)


# =============================================================================
# ❸  The Cleaned-up Web UI (using gr.Blocks)
# =============================================================================
with gr.Blocks(
    title="Python Code Visualizer & Analyzer",
    theme=gr.themes.Soft(primary_hue="blue"),
    css=""" .gradio-container { max-width: 1400px !important; } """,
) as demo:
    # iface_class = gr.Interface(fn=generate_class_diagram_only, inputs=gr.Textbox(), outputs=gr.Image(), api_name="generate_class_diagram", description="Create a UML class diagram (PNG) from Python code.", visible =False)
    # iface_call = gr.Interface(fn=generate_call_graph_only, inputs=gr.Textbox(), outputs=gr.Image(), api_name="generate_call_graph_diagram", description="Generate a function‑call graph (PNG) from Python code.", visible =False)
    # iface_analysis = gr.Interface(fn=analyze_code_only, inputs=gr.Textbox(), outputs=gr.Markdown(), api_name="analyze_code_structure", description="Return a Markdown report with complexity metrics.", visible =False)
    # iface_all = gr.Interface(fn=generate_all_outputs, inputs=gr.Textbox(), outputs=[gr.Image(), gr.Image(), gr.Markdown()], api_name="generate_all", description="Run class diagram, call graph and analysis in one call.", visible =False)
    gr.Markdown(
        """
        # 🐍 Python Code Visualizer & Analyzer
        **Enter Python code, then choose an action to generate diagrams and analysis.**
        This app also functions as an MCP Server, exposing four tools for AI assistants.
        """
    )

    with gr.Row():
        # ---------- Left column – inputs and actions -----------------------------------
        with gr.Column(scale=2):
            gr.Markdown("### 1. Input Code")
            
            example_files = list_example_files()
            print(f"🔍 Found {len(example_files)} example files: {example_files}")
            if example_files:
                example_dropdown = gr.Dropdown(
                    label="Load an Example",
                    choices=example_files,
                    value=example_files[0],
                )
                # initial_code = get_sample_code(example_files[0])
                # initial_code = "# Paste your Python code here\n\nclass MyClass:\n    pass"
                initial_code = "Choose an example file from dropdown or paster your python code here "
                # initial_code = get_sample_code("simple_class.py")
            else:
                initial_code = "# Paste your Python code here\n\nclass MyClass:\n    pass"

            code_input = gr.Textbox(
                label="Python Code",
                placeholder="Paste your Python code here…",
                lines=15,
                max_lines=200,
                value=initial_code,
                elem_classes=["code-input"],
            )

            gr.Markdown("### 2. Choose an Action")
            with gr.Row():
                class_btn = gr.Button("🖼️ Generate Class Diagram")
                call_graph_btn = gr.Button("🔗 Generate Call Graph")
                analyze_btn = gr.Button("📈 Analyze Code")
                all_btn = gr.Button("✨ Generate All", variant="primary")

        # ---------- Right column – outputs ---------------------------------
        with gr.Column(scale=3):
            gr.Markdown("### 3. Results")
            with gr.Tabs():
                with gr.TabItem("UML Class Diagram"):
                    uml_output = gr.Image(label="UML Class Diagram", show_download_button=True, interactive=False)
                with gr.TabItem("Function Call Graph"):
                    call_graph_output = gr.Image(label="Function Call Graph", show_download_button=True, interactive=False)
                with gr.TabItem("Code Analysis Report"):
                    analysis_output = gr.Markdown(label="Comprehensive Code Analysis", elem_classes=["analysis-output"])

    # -------------------------------------------------------------------------
    # Event handlers
    # -------------------------------------------------------------------------

    # Handler to load example code when dropdown changes
    if example_files:
        def _load_example(example_filename: str):
            return get_sample_code(example_filename)
        example_dropdown.change(fn=_load_example, inputs=example_dropdown, outputs=code_input, api_name = False)

    # Handlers for the four action buttons
    # class_btn.click(
    #     fn=generate_class_diagram_only,
    #     inputs=[code_input],
    #     outputs=[uml_output],
    #     api_name=False # Prevents this from creating a duplicate API endpoint
    # )
    class_btn.click(
        fn=generate_class_diagram_only,
        inputs=[code_input],
        outputs=[uml_output],
        # api_name=False # Prevents this from creating a duplicate API endpoint
    )
    
    call_graph_btn.click(
        fn=generate_call_graph_only,
        inputs=[code_input],
        outputs=[call_graph_output],
        # api_name=False
    )

    analyze_btn.click(
        fn=analyze_code_only,
        inputs=[code_input],
        outputs=[analysis_output],
        # api_name=False
    )

    all_btn.click(
        fn=generate_all_outputs,
        inputs=[code_input],
        outputs=[uml_output, call_graph_output, analysis_output],
        # api_name=False
    )

# =============================================================================
# ❹  Launch the App and MCP Server
# =============================================================================
if __name__ == "__main__":
    setup_testing_space()  # Create a persistent working dir if needed

    demo.launch(
        mcp_server=True,    # Enable MCP endpoints (/gradio_api/mcp/*)
        show_api=True,     # Expose ONLY the 4 Interfaces as tools
        show_error=True,    # Display exceptions in the UI
        debug=True,         # Verbose server logs
        share = True,
    )