import gradio as gr
import os
import logging
import gc
import psutil
from functools import wraps
import time
import threading
import json
from model.generate import generate_test_cases, get_generator, monitor_memory

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Thread-safe initialization
_init_lock = threading.Lock()
_initialized = False

def init_model():
    """Initialize model on startup"""
    try:
        # Skip AI model loading in low memory environments
        memory_mb = psutil.Process().memory_info().rss / 1024 / 1024
        if memory_mb > 200 or os.environ.get('HUGGINGFACE_SPACE'):
            logger.info("⚠️ Skipping AI model loading due to memory constraints")
            logger.info("🔧 Using template-based generation mode")
            return True

        logger.info("🚀 Initializing AI model...")
        generator = get_generator()
        model_info = generator.get_model_info()
        logger.info(f"✅ Model initialized: {model_info['model_name']} | Memory: {model_info['memory_usage']}")
        return True
    except Exception as e:
        logger.error(f"❌ Model initialization failed: {e}")
        logger.info("🔧 Falling back to template-based generation")
        return False

def check_health():
    """Check system health"""
    try:
        memory_mb = psutil.Process().memory_info().rss / 1024 / 1024
        return {
            "status": "healthy" if memory_mb < 450 else "warning",
            "memory_usage": f"{memory_mb:.1f}MB",
            "memory_limit": "512MB"
        }
    except Exception:
        return {"status": "unknown", "memory_usage": "unavailable"}

def smart_memory_monitor(func):
    """Enhanced memory monitoring with automatic cleanup"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        try:
            initial_memory = psutil.Process().memory_info().rss / 1024 / 1024
            logger.info(f"🔍 {func.__name__} started | Memory: {initial_memory:.1f}MB")

            if initial_memory > 400:
                logger.warning("⚠️ High memory detected, forcing cleanup...")
                gc.collect()

            result = func(*args, **kwargs)
            return result
        except Exception as e:
            logger.error(f"❌ Error in {func.__name__}: {str(e)}")
            return {
                "error": "Internal server error occurred",
                "message": "Please try again or contact support"
            }
        finally:
            final_memory = psutil.Process().memory_info().rss / 1024 / 1024
            execution_time = time.time() - start_time

            logger.info(f"✅ {func.__name__} completed | Memory: {final_memory:.1f}MB | Time: {execution_time:.2f}s")

            if final_memory > 450:
                logger.warning("🧹 High memory usage, forcing aggressive cleanup...")
                gc.collect()
                post_cleanup_memory = psutil.Process().memory_info().rss / 1024 / 1024
                logger.info(f"🧹 Post-cleanup memory: {post_cleanup_memory:.1f}MB")
    return wrapper

def ensure_initialized():
    """Ensure model is initialized (thread-safe)"""
    global _initialized
    if not _initialized:
        with _init_lock:
            if not _initialized:
                logger.info("🚀 Gradio app starting up on Hugging Face Spaces...")
                success = init_model()
                if success:
                    logger.info("✅ Startup completed successfully")
                else:
                    logger.warning("⚠️ Model initialization failed, using template mode")
                _initialized = True

# Initialize on startup
ensure_initialized()

@smart_memory_monitor
def generate_test_cases_api(srs_text):
    """Main API function for generating test cases"""
    if not srs_text or not srs_text.strip():
        return {
            "error": "No SRS or prompt content provided",
            "test_cases": [],
            "count": 0
        }

    srs_text = srs_text.strip()
    
    if len(srs_text) > 5000:
        logger.warning(f"SRS text truncated from {len(srs_text)} to 5000 characters")
        srs_text = srs_text[:5000]

    try:
        logger.info(f"🎯 Generating test cases for input ({len(srs_text)} chars)")
        test_cases = generate_test_cases(srs_text)

        if not test_cases or len(test_cases) == 0:
            logger.error("No test cases generated")
            return {
                "error": "Failed to generate test cases",
                "test_cases": [],
                "count": 0
            }

        try:
            generator = get_generator()
            model_info = generator.get_model_info()
            model_used = model_info.get("model_name", "Unknown Model")
            generation_method = model_info.get("status", "unknown")
        except Exception:
            model_used = "Template-Based Generator"
            generation_method = "template_mode"

        if model_used == "Template-Based Generator":
            model_algorithm = "Rule-based Template"
            model_reason = "Used rule-based generation due to memory constraints or fallback condition."
        elif "distilgpt2" in model_used:
            model_algorithm = "Transformer-based LM"
            model_reason = "Used DistilGPT2 for balanced performance and memory efficiency."
        elif "DialoGPT" in model_used:
            model_algorithm = "Transformer-based LM"
            model_reason = "Used DialoGPT-small as it fits within memory limits and handles conversational input well."
        else:
            model_algorithm = "Transformer-based LM"
            model_reason = "Used available Hugging Face causal LM due to sufficient resources."

        logger.info(f"✅ Successfully generated {len(test_cases)} test cases")

        return {
            "test_cases": test_cases,
            "count": len(test_cases),
            "model_used": model_used,
            "generation_method": generation_method,
            "model_algorithm": model_algorithm,
            "model_reason": model_reason
        }

    except Exception as e:
        logger.error(f"❌ Test case generation failed: {str(e)}")
        return {
            "error": "Failed to generate test cases",
            "message": "Please try again with different input",
            "test_cases": [],
            "count": 0
        }

def format_test_cases_output(result):
    """Format the test cases for display"""
    if "error" in result:
        return f"❌ Error: {result['error']}", ""
    
    test_cases = result.get("test_cases", [])
    if not test_cases:
        return "No test cases generated", ""
    
    # Format test cases for display
    formatted_output = f"✅ Generated {result['count']} Test Cases\n\n"
    formatted_output += f"🤖 Model: {result['model_used']}\n"
    formatted_output += f"🔧 Algorithm: {result['model_algorithm']}\n"
    formatted_output += f"💡 Reason: {result['model_reason']}\n\n"
    
    formatted_output += "=" * 50 + "\n"
    formatted_output += "TEST CASES:\n"
    formatted_output += "=" * 50 + "\n\n"
    
    for i, tc in enumerate(test_cases, 1):
        formatted_output += f"🔹 Test Case {i}:\n"
        formatted_output += f"   ID: {tc.get('id', f'TC_{i:03d}')}\n"
        formatted_output += f"   Title: {tc.get('title', 'N/A')}\n"
        formatted_output += f"   Description: {tc.get('description', 'N/A')}\n"
        
        steps = tc.get('steps', [])
        if isinstance(steps, list):
            formatted_output += f"   Steps:\n"
            for j, step in enumerate(steps, 1):
                formatted_output += f"      {j}. {step}\n"
        else:
            formatted_output += f"   Steps: {steps}\n"
        
        formatted_output += f"   Expected Result: {tc.get('expected', 'N/A')}\n\n"
    
    # Return JSON for API access
    json_output = json.dumps(result, indent=2)
    
    return formatted_output, json_output

def gradio_generate_test_cases(srs_text):
    """Gradio interface function"""
    result = generate_test_cases_api(srs_text)
    return format_test_cases_output(result)

def get_system_status():
    """Get system status information"""
    health_data = check_health()
    try:
        generator = get_generator()
        model_info = generator.get_model_info()
    except Exception:
        model_info = {
            "model_name": "Template-Based Generator",
            "status": "template_mode",
            "optimization": "memory_safe"
        }

    status_info = f"""
🏥 SYSTEM STATUS
================
Status: {health_data["status"]}
Memory Usage: {health_data["memory_usage"]}
Memory Limit: 512MB

🤖 MODEL INFORMATION
===================
Model Name: {model_info["model_name"]}
Status: {model_info["status"]}
Optimization: {model_info.get("optimization", "standard")}

🚀 APPLICATION INFO
==================
Version: 1.0.0-spaces-optimized
Environment: Hugging Face Spaces
Backend: Gradio
    """
    return status_info

def get_model_info_detailed():
    """Get detailed model information"""
    try:
        generator = get_generator()
        info = generator.get_model_info()
        health_data = check_health()

        detailed_info = f"""
🔧 DETAILED MODEL INFORMATION
============================
Model Name: {info.get('model_name', 'N/A')}
Status: {info.get('status', 'N/A')}
Memory Usage: {info.get('memory_usage', 'N/A')}
Optimization Level: {info.get('optimization', 'N/A')}

📊 SYSTEM METRICS
================
System Status: {health_data['status']}
Current Memory: {health_data['memory_usage']}
Memory Limit: {health_data.get('memory_limit', 'N/A')}

⚙️ CONFIGURATION
===============
Environment: {"Hugging Face Spaces" if os.environ.get('SPACE_ID') else "Local"}
Backend: Gradio
Threading: Enabled
Memory Monitoring: Active
        """
        return detailed_info
    except Exception as e:
        return f"❌ Error getting model info: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="AI Test Case Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # 🧪 AI Test Case Generator
    
    Generate comprehensive test cases from Software Requirements Specification (SRS) documents using AI models.
    
    **Features:**
    - 🤖 AI-powered test case generation
    - 📝 Template-based fallback for low memory environments  
    - 🔧 Memory-optimized processing
    - 📊 Real-time system monitoring
    """)
    
    with gr.Tab("🧪 Generate Test Cases"):
        with gr.Row():
            with gr.Column(scale=2):
                srs_input = gr.Textbox(
                    label="📋 Software Requirements Specification (SRS)",
                    placeholder="Enter your SRS document or requirements here...\n\nExample:\nThe system shall provide user authentication with username and password. Users should be able to login, logout, and reset passwords. The system should validate input and display appropriate error messages for invalid credentials.",
                    lines=10,
                    max_lines=20
                )
                
                generate_btn = gr.Button("🚀 Generate Test Cases", variant="primary", size="lg")
            
            with gr.Column(scale=3):
                output_display = gr.Textbox(
                    label="📋 Generated Test Cases",
                    lines=20,
                    max_lines=30,
                    interactive=False
                )
        
        with gr.Row():
            json_output = gr.Textbox(
                label="📄 JSON Output (for API use)",
                lines=10,
                max_lines=15,
                interactive=False
            )
    
    with gr.Tab("📊 System Status"):
        with gr.Column():
            status_display = gr.Textbox(
                label="🏥 System Health & Status",
                lines=15,
                interactive=False
            )
            refresh_status_btn = gr.Button("🔄 Refresh Status", variant="secondary")
    
    with gr.Tab("🔧 Model Information"):
        with gr.Column():
            model_info_display = gr.Textbox(
                label="🤖 Detailed Model Information",
                lines=20,
                interactive=False
            )
            refresh_model_btn = gr.Button("🔄 Refresh Model Info", variant="secondary")
    
    with gr.Tab("📚 API Documentation"):
        gr.Markdown("""
        ## 🔌 API Endpoints
        
        This Gradio app automatically creates API endpoints for all functions:
        
        ### Generate Test Cases
        **Endpoint:** `/api/predict`  
        **Method:** POST  
        **Body:**
        ```json
        {
          "data": ["Your SRS text here"]
        }
        ```
        
        **Response:**
        ```json
        {
          "data": [
            "Formatted test cases output",
            "JSON output with test cases"
          ]
        }
        ```
        
        ### Example Usage (Python):
        ```python
        import requests
        
        response = requests.post(
            "YOUR_SPACE_URL/api/predict",
            json={"data": ["User login system requirements..."]}
        )
        result = response.json()
        test_cases = result["data"][1]  # JSON output
        ```
        
        ### Example Usage (cURL):
        ```bash
        curl -X POST "YOUR_SPACE_URL/api/predict" \\
          -H "Content-Type: application/json" \\
          -d '{"data":["Your SRS text here"]}'
        ```
        
        ## 📋 Response Format
        
        The API returns test cases in this format:
        ```json
        {
          "test_cases": [
            {
              "id": "TC_001",
              "title": "Test Case Title", 
              "description": "Test description",
              "steps": ["Step 1", "Step 2"],
              "expected": "Expected result"
            }
          ],
          "count": 3,
          "model_used": "distilgpt2",
          "model_algorithm": "Transformer-based LM",
          "model_reason": "Selected for balanced performance..."
        }
        ```
        """)
    
    # Event handlers
    generate_btn.click(
        fn=gradio_generate_test_cases,
        inputs=[srs_input],
        outputs=[output_display, json_output]
    )
    
    refresh_status_btn.click(
        fn=get_system_status,
        outputs=[status_display]
    )
    
    refresh_model_btn.click(
        fn=get_model_info_detailed,
        outputs=[model_info_display]
    )
    
    # Load initial status
    app.load(
        fn=get_system_status,
        outputs=[status_display]
    )
    
    app.load(
        fn=get_model_info_detailed,
        outputs=[model_info_display]
    )

# Launch the app
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    
    logger.info(f"🚀 Starting Gradio app on port {port}")
    logger.info(f"🖥️ Environment: {'Hugging Face Spaces' if os.environ.get('SPACE_ID') else 'Local'}")
    
    app.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False,
        show_error=True
    )