"""
Self-Documentation Agent System - Gradio Demo
A Hugging Face Spaces demo for the LangGraph-based code documentation workflow
"""

import os
import gradio as gr
import tempfile
import zipfile
from io import BytesIO
from typing import Dict, Tuple
import datetime

# Import the main workflow function
from main import run_documentation_workflow

def process_code(code_input: str, google_api_key: str, tavily_api_key: str) -> Tuple[str, str, str]:
    """
    Process the input code through the documentation workflow
    
    Args:
        code_input: Python code to analyze and document
        google_api_key: Google API key for Gemini model
        tavily_api_key: Tavily API key for web search
        
    Returns:
        Tuple of (documented_code, analysis_results, status_message)
    """
    if not code_input.strip():
        return "", "", "❌ Please provide Python code to analyze"
    
    if not google_api_key.strip():
        return "", "", "❌ Please provide a Google API key. Get one from [Google AI Studio](https://makersuite.google.com/app/apikey)"
    
    if not tavily_api_key.strip():
        return "", "", "❌ Please provide a Tavily API key. Get one from [Tavily](https://tavily.com/)"
    
    try:
        # Set environment variables
        os.environ["GOOGLE_API_KEY"] = google_api_key.strip()
        os.environ["TAVILY_API_KEY"] = tavily_api_key.strip()
        
        # Show processing status
        status_msg = "🔄 Processing your code through the documentation workflow..."
        
        # Run the workflow
        result = run_documentation_workflow(code_input)
        
        # Read the generated files
        documented_code = ""
        analysis_results = ""
        
        try:
            with open("code.py", "r", encoding="utf-8") as f:
                documented_code = f.read()
        except FileNotFoundError:
            documented_code = "# No documented code generated\n# The workflow may have encountered an issue"
        
        try:
            with open("analysis.txt", "r", encoding="utf-8") as f:
                analysis_results = f.read()
        except FileNotFoundError:
            analysis_results = "No analysis results generated. The workflow may have encountered an issue."
        
        # Create status message
        libraries_count = len(result.get('libraries_used', []))
        issues_count = len(result.get('issues_found', []))
        tests_count = len(result.get('test_results', []))
        
        status = f"""✅ **Workflow completed successfully!**

📊 **Summary:**
- 📚 Libraries found: {libraries_count}
- ⚠️ Issues identified: {issues_count}
- 🧪 Test results: {tests_count}
- 📝 Status: {result.get('current_step', 'completed')}

💡 **Next Steps:**
1. Review the documented code on the left
2. Check the analysis results on the right
3. Address any issues or recommendations
4. Test your code with the suggested improvements
        """
        
        return documented_code, analysis_results, status
        
    except Exception as e:
        error_details = str(e)
        if "API key" in error_details.lower():
            error_msg = f"❌ **API Key Error**: {error_details}\n\nPlease check that your API keys are valid and have the necessary permissions."
        elif "quota" in error_details.lower() or "limit" in error_details.lower():
            error_msg = f"❌ **Rate Limit Error**: {error_details}\n\nYou may have exceeded your API quota. Please try again later."
        else:
            error_msg = f"❌ **Processing Error**: {error_details}\n\nPlease check your code syntax and try again. If the issue persists, it may be a temporary service problem."
        
        return "", "", error_msg

def create_download_file() -> str:
    """Create a zip file with the generated outputs for download"""
    try:
        # Create a temporary zip file
        zip_buffer = BytesIO()
        
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            # Add code.py if it exists
            if os.path.exists("code.py"):
                zip_file.write("code.py", "documented_code.py")
            
            # Add analysis.txt if it exists
            if os.path.exists("analysis.txt"):
                zip_file.write("analysis.txt", "analysis_results.txt")
            
            # Add workflow diagram if it exists
            if os.path.exists("workflow_diagram.png"):
                zip_file.write("workflow_diagram.png", "workflow_diagram.png")
        
        # Save to temporary file
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.zip')
        temp_file.write(zip_buffer.getvalue())
        temp_file.close()
        
        return temp_file.name
        
    except Exception as e:
        print(f"Error creating download file: {e}")
        return None

# Sample code for demonstration
SAMPLE_CODE = '''import math
import random

def calculate_area(shape, **kwargs):
    if shape == "circle":
        return math.pi * kwargs["radius"] ** 2
    elif shape == "rectangle":
        return kwargs["width"] * kwargs["height"]
    else:
        return 0

def divide_numbers(a, b):
    return a / b

def process_list(items):
    total = 0
    for i in range(len(items)):
        total += items[i] * 2
    return total

class Calculator:
    def __init__(self):
        self.history = []

    def add(self, a, b):
        result = a + b
        self.history.append(f"{a} + {b} = {result}")
        return result

    def divide(self, a, b):
        return divide_numbers(a, b)

calc = Calculator()
result = calc.add(5, 3)
area = calculate_area("circle", radius=5)
division = calc.divide(10, 2)
items = [1, 2, 3, 4]
processed = process_list(items)
print(f"Results: {result}, {area:.2f}, {division}, {processed}")'''

# Create the Gradio interface
def create_interface():
    """Create the Gradio interface for the documentation system"""
    
    with gr.Blocks(
        title="Self-Documentation Agent System",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .code-output {
            font-family: 'Courier New', monospace;
            font-size: 12px;
        }
        """
    ) as demo:
        
        gr.Markdown("""
        # 🤖 Self-Documentation Agent System
        
        An AI-powered system that automatically analyzes, documents, and validates Python code using LangGraph workflows and Google's Gemini 2.5 Flash model.
        
        ## ✨ Features
        - **🧠 Intelligent Documentation**: Adds comprehensive docstrings and comments
        - **🔍 Code Analysis**: Identifies issues and provides recommendations  
        - **🧪 Automated Testing**: Executes code with various test scenarios
        - **📋 Professional Output**: Clean, business-appropriate formatting
        - **🌐 Library Research**: Uses web search to understand unfamiliar libraries
        
        ## 🚀 How to Use
        1. **🔑 Enter API Keys**: Get your free keys from the links below
        2. **📝 Add Your Code**: Paste Python code or use the sample
        3. **⚡ Process**: Click "Analyze & Document Code"
        4. **📊 Review Results**: Check documented code and analysis
        
        > **💡 Tip**: The system works best with functions, classes, and complete code snippets!
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 🔑 API Configuration")
                google_key = gr.Textbox(
                    label="🔑 Google API Key",
                    placeholder="Enter your Google API key for Gemini access",
                    type="password",
                    info="🔗 Get your free key: https://makersuite.google.com/app/apikey"
                )
                tavily_key = gr.Textbox(
                    label="🔍 Tavily API Key", 
                    placeholder="Enter your Tavily API key for web search",
                    type="password",
                    info="🔗 Get your free key: https://tavily.com/"
                )
                
                gr.Markdown("### 📝 Input Code")
                code_input = gr.Code(
                    label="Python Code to Analyze",
                    language="python",
                    value=SAMPLE_CODE,
                    lines=20
                )
                
                analyze_btn = gr.Button(
                    "🚀 Analyze & Document Code",
                    variant="primary",
                    size="lg"
                )
                
                # Add sample code button
                sample_btn = gr.Button(
                    "📋 Load Sample Code",
                    variant="secondary"
                )
        
        with gr.Row():
            status_output = gr.Markdown(
                label="Status",
                value="Ready to analyze code. Please provide your API keys and code."
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📚 Documented Code")
                documented_output = gr.Code(
                    label="Generated Documentation",
                    language="python",
                    lines=25,
                    elem_classes=["code-output"]
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### 🔍 Analysis Results")
                analysis_output = gr.Textbox(
                    label="Code Analysis & Test Results",
                    lines=25,
                    elem_classes=["code-output"],
                    info="Detailed analysis, issues, and test results"
                )
        
        # Event handlers
        analyze_btn.click(
            fn=process_code,
            inputs=[code_input, google_key, tavily_key],
            outputs=[documented_output, analysis_output, status_output]
        )
        
        sample_btn.click(
            fn=lambda: SAMPLE_CODE,
            outputs=code_input
        )
        
        # Add information section
        with gr.Accordion("ℹ️ About the Workflow", open=False):
            gr.Markdown("""
            ### Workflow Steps
            
            1. **Research Node**: Analyzes code structure, identifies libraries, and checks existing documentation
            2. **Document Node**: Adds comprehensive docstrings and inline comments using Google style
            3. **Analyze Node**: Executes code with various test scenarios and identifies potential issues
            4. **Final Node**: Saves results and provides summary
            
            ### What You Get
            
            - **Enhanced Code**: Your original code with professional documentation
            - **Detailed Analysis**: Comprehensive testing results and recommendations
            - **Issue Identification**: Potential problems and improvement suggestions
            - **Usage Guidelines**: Instructions for implementation and testing
            
            ### API Keys Required
            
            - **Google API Key**: Get it from [Google AI Studio](https://makersuite.google.com/app/apikey)
            - **Tavily API Key**: Get it from [Tavily](https://tavily.com/) for web search functionality
            """)
        
        with gr.Accordion("🔧 Technical Details", open=False):
            gr.Markdown("""
            ### Technologies Used
            
            - **LangGraph**: Workflow orchestration and state management
            - **Google Gemini 2.5 Flash**: Advanced language model for analysis
            - **Tavily Search**: Web search for library research
            - **Python AST**: Code structure analysis
            - **Code Execution Tools**: Safe code testing environment
            
            ### Supported Features
            
            - Function and class documentation
            - Inline code comments
            - Error handling analysis
            - Performance recommendations
            - Security issue identification
            - Test case generation
            - I/O behavior documentation
            """)
    
    return demo

# Create and launch the interface
if __name__ == "__main__":
    demo = create_interface()
    demo.launch()