app.py

import os
import gradio as gr
import json
import plotly.express as px
import pandas as pd
from typing import Tuple, Dict, Optional, List
import logging
logger = logging.getLogger(__name__)

# Import from our modules
from utils import (
    validate_apex_syntax, perform_skeptical_evaluation, extract_code_blocks,
    format_structured_explanation, format_object_conversion_explanation,
    extract_validation_metrics, normalize_metrics, generate_test_cases,
    VALIDATION_SCHEMA, B2B_COMMERCE_PATTERNS, logger
)

from api_client import (
    all_models, together_models, anthropic_models, call_llm
)

def correct_apex_trigger(model: str, trigger_code: str, progress=None) -> Tuple[str, str, str]:
    """Correct Apex Trigger with skeptical evaluation."""
    if progress:
        progress(0.1, desc="Validating input...")

    # Input validation
    if not trigger_code.strip():
        return "Please provide Apex Trigger code to correct.", "", ""

    if len(trigger_code.strip()) < 50:
        return "Code too short to be a valid Apex trigger. Please provide complete code.", "", ""

    # Perform initial syntax check
    is_valid, syntax_issues = validate_apex_syntax(trigger_code)

    if progress:
        progress(0.3, desc="Analyzing code structure...")

    # Perform skeptical evaluation
    evaluation = perform_skeptical_evaluation(trigger_code, "trigger")

    # Build comprehensive prompt with structured output format
    prompt = f"""
Correct this Apex Trigger for B2B Lightning Experience migration. Be BRIEF and DIRECT.
ORIGINAL CODE:
```apex
{trigger_code}
```
DETECTED ISSUES:
- Syntax Errors: {len(syntax_issues)}
- Security Issues: {len(evaluation['security_concerns'])}
- Performance Issues: {len(evaluation['performance_issues'])}
- B2B Commerce Issues: {len(evaluation['b2b_commerce_issues'])}
PROVIDE EXACTLY THIS FORMAT:
## CORRECTED CODE
```apex
[Put the complete corrected trigger here with inline comments for changes]
```
## KEY CHANGES (bullet points only)
- [Change 1: Brief description]
- [Change 2: Brief description]
- [Maximum 7 bullet points]
## CRITICAL ISSUES FIXED
1. [Most important issue]: [One-line explanation]
2. [Second issue]: [One-line explanation]
3. [Third issue]: [One-line explanation]
## REMAINING WARNINGS
- [Any issues that need manual review]
BE CONCISE. NO VERBOSE EXPLANATIONS. FOCUS ON CODE QUALITY.
"""

    if progress:
        progress(0.5, desc="Calling AI model...")

    response = call_llm(model, prompt, temperature=0.2)  # Even lower temperature for consistency

    if progress:
        progress(0.8, desc="Processing response...")

    # Extract code and explanations
    code_output = extract_code_blocks(response)

    # Validate the corrected code
    validation_warnings = []
    if code_output:
        corrected_valid, corrected_issues = validate_apex_syntax(code_output)
        if not corrected_valid:
            error_count = len([i for i in corrected_issues if i["type"] == "error"])
            warning_count = len([i for i in corrected_issues if i["type"] == "warning"])
            validation_warnings.append(f"⚠️ Validation: {error_count} errors, {warning_count} warnings remain")

    # Extract structured explanation
    explanation = format_structured_explanation(response, code_output)

    # Add validation warnings if any
    if validation_warnings:
        explanation = "\n".join(validation_warnings) + "\n\n" + explanation

    # Add test case reference (brief)
    explanation += "\n\n**Test Template:** Available in Full Response section"

    if progress:
        progress(1.0, desc="Complete!")

    return response, code_output, explanation

def convert_cc_object(model: str, cc_object_code: str, progress=None) -> Tuple[str, str, str]:
    """Convert CloudCraze Object with skeptical evaluation."""
    if progress:
        progress(0.1, desc="Validating input...")

    # Input validation
    if not cc_object_code.strip():
        return "Please provide CloudCraze Object code to convert.", "", ""

    if len(cc_object_code.strip()) < 30:
        return "Code too short to be a valid CloudCraze object. Please provide complete code.", "", ""

    if progress:
        progress(0.3, desc="Analyzing CloudCraze structure...")

    # Check for CloudCraze patterns
    import re
    has_cc_pattern = bool(re.search(B2B_COMMERCE_PATTERNS["cloudcraze_reference"], cc_object_code))
    if not has_cc_pattern:
        logger.warning("No obvious CloudCraze patterns found in input")

    # Perform evaluation
    evaluation = perform_skeptical_evaluation(cc_object_code, "object")

    prompt = f"""
Convert this CloudCraze Object to B2B Lightning Experience. Be BRIEF and PRECISE.
CLOUDCRAZE OBJECT:
```
{cc_object_code}
```
PROVIDE EXACTLY THIS FORMAT:
## B2B LEX OBJECT MAPPING
- Source: [CloudCraze object name]
- Target: [B2B LEX object name]
- Migration Complexity: [Low/Medium/High]
## CONVERTED CODE
```apex
[Complete B2B LEX implementation with inline comments]
```
## FIELD MAPPINGS (table format)
| CC Field | B2B Field | Type | Notes |
|----------|-----------|------|-------|
| field1   | newField1 | Text | Required |
| field2   | newField2 | Number | Optional |
## MIGRATION STEPS
1. [Step 1 - one line]
2. [Step 2 - one line]
3. [Maximum 5 steps]
## DATA MIGRATION SCRIPT
```apex
[Brief data migration code if needed, otherwise state "Not Required"]
```
## WARNINGS
- [Any manual steps or considerations]
BE CONCISE. FOCUS ON ACTIONABLE INFORMATION.
"""

    if progress:
        progress(0.5, desc="Calling AI model...")

    response = call_llm(model, prompt, temperature=0.2)

    if progress:
        progress(0.8, desc="Processing response...")

    # Extract code and explanations
    code_output = extract_code_blocks(response)

    # Extract structured explanation
    explanation = format_object_conversion_explanation(response, code_output)

    # Add brief test reference
    explanation += "\n\n**Test Utilities:** See Full Response for test data creation utilities"

    if progress:
        progress(1.0, desc="Complete!")

    return response, code_output, explanation

def validate_apex_trigger(validation_model: str, original_code: str, corrected_code: str) -> str:
    """Enhanced validation with skeptical evaluation and structured output."""
    if not validation_model or not original_code.strip() or not corrected_code.strip():
        return "Please provide all required inputs for validation."

    # Perform syntax validation on both
    orig_valid, orig_issues = validate_apex_syntax(original_code)
    corr_valid, corr_issues = validate_apex_syntax(corrected_code)

    prompt = f"""
Validate this Apex trigger correction. Be CRITICAL but CONCISE.
ORIGINAL ({len(orig_issues)} issues detected):
```apex
{original_code}
```
CORRECTED ({len(corr_issues)} issues detected):
```apex
{corrected_code}
```
PROVIDE THIS EXACT JSON FORMAT:
```json
{{
    "quality_rating": [1-10],
    "accuracy": [0.0-1.0],
    "completeness": [0.0-1.0],
    "best_practices_alignment": [0.0-1.0],
    "syntax_validity": [0.0-1.0],
    "security_score": [0.0-1.0],
    "performance_score": [0.0-1.0],
    "errors": [
        "Error 1: [specific line/issue]",
        "Error 2: [specific line/issue]"
    ],
    "warnings": [
        "Warning 1: [potential issue]",
        "Warning 2: [potential issue]"
    ],
    "suggestions": [
        "Improvement 1: [actionable suggestion]",
        "Improvement 2: [actionable suggestion]"
    ]
}}
``` 
SCORING CRITERIA:
- quality_rating: Overall code quality (1=terrible, 10=excellent)
- accuracy: Correctness of fixes (0=wrong, 1=perfect)
- completeness: All issues addressed (0=none, 1=all)
- best_practices_alignment: Follows Salesforce standards (0=poor, 1=excellent)
- syntax_validity: No syntax errors (0=many errors, 1=error-free)
- security_score: Security best practices (0=vulnerable, 1=secure)
- performance_score: Efficiency and scalability (0=poor, 1=optimal)
BE HARSH. FIND PROBLEMS. Maximum 3 items per category.
"""

    return call_llm(validation_model, prompt, temperature=0.1)  # Very low temperature for consistent JSON

def validate_cc_object_conversion(validation_model: str, original_object: str, converted_object: str) -> str:
    """Enhanced validation for CloudCraze object conversion with structured output."""
    if not validation_model or not original_object.strip() or not converted_object.strip():
        return "Please provide all required inputs for validation."

    prompt = f"""
Validate this CloudCraze to B2B LEX conversion. Be CRITICAL and BRIEF.
ORIGINAL CLOUDCRAZE:
```
{original_object}
```
CONVERTED B2B LEX:
```
{converted_object}
```
PROVIDE THIS EXACT JSON FORMAT:
```json
{{
    "quality_rating": [1-10],
    "accuracy": [0.0-1.0],
    "completeness": [0.0-1.0],
    "best_practices_alignment": [0.0-1.0],
    "syntax_validity": [0.0-1.0],
    "security_score": [0.0-1.0],
    "performance_score": [0.0-1.0],
    "errors": [
        "Missing field: [field_name]",
        "Wrong mapping: [issue]"
    ],
    "warnings": [
        "Data type mismatch: [field]",
        "Custom logic not migrated: [what]"
    ],
    "suggestions": [
        "Add validation for: [field]",
        "Consider indexing: [field]"
    ]
}}
```
FOCUS ON:
1. Missing field mappings
2. Data type conversions
3. Relationship integrity
4. Custom field handling
5. Performance at scale
BE HARSH. Maximum 3 items per category. Focus on REAL issues.
"""

    return call_llm(validation_model, prompt, temperature=0.1)

def create_enhanced_radar_chart(metrics: Optional[Dict[str, float]]) -> Optional[object]:
    """Create an enhanced radar chart with more metrics."""
    if not metrics:
        return None

    # Create data for the radar chart
    categories = [
        "Quality",
        "Accuracy",
        "Completeness",
        "Best Practices",
        "Syntax Valid",
        "Security",
        "Performance"
    ]

    values = [
        metrics.get("quality_rating", 0) / 10,  # Normalize to 0-1 scale
        metrics.get("accuracy", 0),
        metrics.get("completeness", 0),
        metrics.get("best_practices_alignment", 0),
        metrics.get("syntax_validity", 0),
        metrics.get("security_score", 0),
        metrics.get("performance_score", 0)
    ]

    # Create a DataFrame for plotting
    df = pd.DataFrame({
        'Category': categories,
        'Value': values
    })

    # Create the radar chart
    fig = px.line_polar(
        df, r='Value', theta='Category', line_close=True,
        range_r=[0, 1], title="Comprehensive Validation Assessment"
    )
    fig.update_traces(fill='toself', fillcolor='rgba(0, 100, 255, 0.2)')
    fig.update_layout(
        polar=dict(
            radialaxis=dict(
                visible=True,
                range=[0, 1]
            )
        ),
        showlegend=False,
        height=400
    )

    return fig

def get_theme_styles(theme_choice: str) -> Tuple[str, str, str, str]:
    """Get theme styles for different UI elements."""
    themes = {
        "Dark": {
            "explanation": "background-color: #1e1e1e; color: #e0e0e0; padding: 15px; border-radius: 8px; font-family: 'Inter', sans-serif; line-height: 1.6;",
            "code": "background-color: #0d1117; color: #c9d1d9; font-family: 'Fira Code', 'Consolas', monospace; padding: 15px; border-radius: 8px; border: 1px solid #30363d; font-size: 14px; line-height: 1.5;",
            "validation": "background-color: #161b22; color: #c9d1d9; padding: 15px; border-radius: 8px; border: 1px solid #30363d;",
            "error": "background-color: #3d1418; color: #f85149; padding: 10px; border-radius: 5px; border: 1px solid #f85149;"
        },
        "Light": {
            "explanation": "background-color: #ffffff; color: #24292e; padding: 15px; border-radius: 8px; border: 1px solid #e1e4e8; font-family: 'Inter', sans-serif; line-height: 1.6;",
            "code": "background-color: #f6f8fa; color: #24292e; font-family: 'Fira Code', 'Consolas', monospace; padding: 15px; border-radius: 8px; border: 1px solid #e1e4e8; font-size: 14px; line-height: 1.5;",
            "validation": "background-color: #f6f8fa; color: #24292e; padding: 15px; border-radius: 8px; border: 1px solid #e1e4e8;",
            "error": "background-color: #ffe4e6; color: #d73a49; padding: 10px; border-radius: 5px; border: 1px solid #d73a49;"
        }
    }

    theme = themes.get(theme_choice, themes["Light"])
    return (theme["explanation"], theme["code"], theme["explanation"], theme["code"])

# Wrapper functions with proper error handling
def trigger_correction_wrapper(model, code):
    """Wrapper for trigger correction with proper error handling."""
    try:
        if not model or not code.strip():
            return "Please select a model and provide code.", "", "Please provide valid inputs."

        # Create dummy progress function since Gradio progress doesn't work in lambda
        def dummy_progress(value, desc=""):
            pass

        return correct_apex_trigger(model, code, progress=dummy_progress)
    except Exception as e:
        logger.error(f"Trigger correction error: {str(e)}")
        error_msg = f"Error processing request: {str(e)}"
        return error_msg, "", error_msg

def object_conversion_wrapper(model, code):
    """Wrapper for object conversion with proper error handling."""
    try:
        if not model or not code.strip():
            return "Please select a model and provide code.", "", "Please provide valid inputs."

        def dummy_progress(value, desc=""):
            pass

        return convert_cc_object(model, code, progress=dummy_progress)
    except Exception as e:
        logger.error(f"Object conversion error: {str(e)}")
        error_msg = f"Error processing request: {str(e)}"
        return error_msg, "", error_msg

def validate_and_chart_trigger(model, original, corrected):
    """Wrapper for trigger validation with error handling."""
    try:
        if not model or not original.strip() or not corrected.strip():
            return "Please provide all required inputs for validation.", None

        validation_text = validate_apex_trigger(model, original, corrected)
        metrics = extract_validation_metrics(validation_text)
        chart = create_enhanced_radar_chart(metrics) if metrics else None
        return validation_text, chart
    except Exception as e:
        logger.error(f"Trigger validation error: {str(e)}")
        return f"Validation error: {str(e)}", None

def validate_and_chart_object(model, original, converted):
    """Wrapper for object validation with error handling."""
    try:
        if not model or not original.strip() or not converted.strip():
            return "Please provide all required inputs for validation.", None

        validation_text = validate_cc_object_conversion(model, original, converted)
        metrics = extract_validation_metrics(validation_text)
        chart = create_enhanced_radar_chart(metrics) if metrics else None
        return validation_text, chart
    except Exception as e:
        logger.error(f"Object validation error: {str(e)}")
        return f"Validation error: {str(e)}", None

def main():
    """Main application entry point."""
    with gr.Blocks(
            title="Salesforce B2B Commerce Migration Assistant",
            theme=gr.themes.Soft(primary_hue="blue"),
            css="""
        .gradio-container {
            font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
        }
        .gr-button-primary {
            background-color: #0969da !important;
        }
        .gr-button-primary:hover {
            background-color: #0860ca !important;
        }
        code {
            font-family: 'Fira Code', 'Consolas', 'Monaco', monospace;
        }
        """
    ) as app:

        gr.Markdown("# 🚀 Salesforce B2B Commerce Migration Assistant")
        gr.Markdown("Advanced tool for migrating CloudCraze code to B2B Lightning Experience with skeptical AI validation.")

        # Model selection
        with gr.Row():
            with gr.Column():
                gr.Markdown("### 🤖 Primary Model")
                primary_model_dropdown = gr.Dropdown(
                    choices=all_models,
                    value=anthropic_models[0],
                    label="Select Primary AI Model for Conversion",
                    info="This model performs the initial code conversion"
                )

            with gr.Column():
                gr.Markdown("### 🔍 Validation Model")
                validation_model_dropdown = gr.Dropdown(
                    choices=all_models,
                    value=anthropic_models[1],
                    label="Select Validation AI Model for Review",
                    info="This model skeptically reviews and validates the output"
                )

        with gr.Tab("⚡ Apex Trigger Correction"):
            gr.Markdown("### Apex Trigger Correction & Optimization")
            gr.Markdown("Paste your Apex Trigger code for analysis, correction, and optimization.")

            trigger_input = gr.Textbox(
                lines=15,
                placeholder="Paste your Apex Trigger code here...\n\nExample:\ntrigger AccountTrigger on Account (before insert, before update) {\n    // Your trigger logic\n}",
                label="Apex Trigger Code",
                elem_classes="code-input"
            )

            with gr.Row():
                trigger_button = gr.Button("🔧 Correct & Optimize", variant="primary", size="lg")
                copy_code_button = gr.Button("📋 Copy Code", variant="secondary")

            # Progress indicator
            trigger_progress = gr.Textbox(label="Progress", visible=False)

            with gr.Accordion("📄 Full Model Response", open=False):
                trigger_full_response = gr.Textbox(
                    lines=20,
                    label="Full Model Response",
                    interactive=False
                )

            with gr.Row():
                with gr.Column():
                    trigger_explanation = gr.Textbox(
                        lines=15,
                        label="📝 Explanation & Analysis",
                        placeholder="Detailed explanation will appear here...",
                        interactive=False,
                        elem_id="trigger_explanation"
                    )

                with gr.Column():
                    trigger_code_output = gr.Code(
                        language="python",  # Using python highlighting as java is not supported
                        label="✅ Corrected Code (Apex)",
                        value="// Corrected Apex code will appear here",
                        elem_id="trigger_code_output"
                    )

            gr.Markdown("### 🎯 Validation Results")
            with gr.Row():
                with gr.Column(scale=2):
                    trigger_validation_output = gr.Textbox(
                        lines=20,
                        label="🔍 Skeptical Validation Assessment",
                        placeholder="Validation results will appear here...",
                        interactive=True,
                        elem_id="trigger_validation"
                    )
                with gr.Column(scale=1):
                    trigger_chart = gr.Plot(label="📊 Validation Metrics")

            validate_trigger_button = gr.Button("🔍 Validate Correction", variant="secondary", size="lg")

            # Wire up functionality - INSIDE the main function where UI elements are defined
            trigger_button.click(
                fn=trigger_correction_wrapper,
                inputs=[primary_model_dropdown, trigger_input],
                outputs=[trigger_full_response, trigger_code_output, trigger_explanation],
                show_progress=True
            )

            validate_trigger_button.click(
                fn=validate_and_chart_trigger,
                inputs=[validation_model_dropdown, trigger_input, trigger_code_output],
                outputs=[trigger_validation_output, trigger_chart],
                show_progress=True
            )

            copy_code_button.click(
                fn=lambda: gr.Info("Code copied! Use Ctrl+C/Cmd+C if automatic copy fails."),
                inputs=[],
                outputs=[]
            )

            with gr.Row():
                trigger_clear = gr.Button("🗑️ Clear Input")
                trigger_clear.click(lambda: "", [], trigger_input)

                results_clear = gr.Button("🧹 Clear Results")
                results_clear.click(
                    lambda: ["", "", "", "", None],
                    [],
                    [trigger_full_response, trigger_code_output, trigger_explanation, trigger_validation_output, trigger_chart]
                )

        with gr.Tab("🔄 CloudCraze Object Conversion"):
            gr.Markdown("### CloudCraze to B2B Lightning Experience Object Conversion")
            gr.Markdown("Convert CloudCraze custom objects to B2B Lightning Experience format.")

            object_input = gr.Textbox(
                lines=15,
                placeholder="Paste your CloudCraze Object definition here...\n\nExample:\nE_Product__c fields, relationships, and custom logic",
                label="CloudCraze Object Code",
                elem_classes="code-input"
            )

            with gr.Row():
                object_button = gr.Button("🔄 Convert Object", variant="primary", size="lg")
                object_copy_code_button = gr.Button("📋 Copy Code", variant="secondary")

            # Progress indicator
            object_progress = gr.Textbox(label="Progress", visible=False)

            with gr.Accordion("📄 Full Model Response", open=False):
                object_full_response = gr.Textbox(
                    lines=20,
                    label="Full Model Response",
                    interactive=False
                )

            with gr.Row():
                with gr.Column():
                    object_explanation = gr.Textbox(
                        lines=15,
                        label="📝 Conversion Explanation",
                        placeholder="Detailed explanation will appear here...",
                        interactive=False,
                        elem_id="object_explanation"
                    )

                with gr.Column():
                    object_code_output = gr.Code(
                        language="python",  # Using python highlighting as java is not supported
                        label="✅ Converted Code (B2B LEX)",
                        value="// Converted B2B Lightning Experience code will appear here",
                        elem_id="object_code_output"
                    )

            gr.Markdown("### 🎯 Validation Results")
            with gr.Row():
                with gr.Column(scale=2):
                    object_validation_output = gr.Textbox(
                        lines=20,
                        label="🔍 Skeptical Validation Assessment",
                        placeholder="Validation results will appear here...",
                        interactive=True,
                        elem_id="object_validation"
                    )
                with gr.Column(scale=1):
                    object_chart = gr.Plot(label="📊 Validation Metrics")

            validate_object_button = gr.Button("🔍 Validate Conversion", variant="secondary", size="lg")

            # Wire up functionality - INSIDE the main function where UI elements are defined
            object_button.click(
                fn=object_conversion_wrapper,
                inputs=[primary_model_dropdown, object_input],
                outputs=[object_full_response, object_code_output, object_explanation],
                show_progress=True
            )

            validate_object_button.click(
                fn=validate_and_chart_object,
                inputs=[validation_model_dropdown, object_input, object_code_output],
                outputs=[object_validation_output, object_chart],
                show_progress=True
            )

            object_copy_code_button.click(
                fn=lambda: gr.Info("Code copied! Use Ctrl+C/Cmd+C if automatic copy fails."),
                inputs=[],
                outputs=[]
            )

            with gr.Row():
                object_clear = gr.Button("🗑️ Clear Input")
                object_clear.click(lambda: "", [], object_input)

                object_results_clear = gr.Button("🧹 Clear Results")
                object_results_clear.click(
                    lambda: ["", "", "", "", None],
                    [],
                    [object_full_response, object_code_output, object_explanation, object_validation_output, object_chart]
                )

        # UI Preferences
        with gr.Accordion("⚙️ UI Preferences", open=False):
            theme_radio = gr.Radio(
                label="🎨 Theme",
                choices=["Light", "Dark"],
                value="Light"
            )

            # Comment out theme functionality for now to avoid conflicts
            # theme_radio.change(
            #     fn=get_theme_styles,
            #     inputs=[theme_radio],
            #     outputs=[
            #         trigger_explanation,
            #         trigger_code_output,
            #         object_explanation,
            #         object_code_output
            #     ]
            # )

        gr.Markdown("### 📚 About This Tool")
        gr.Markdown(
            """
            **🚀 Enhanced Features:**
            - **Skeptical AI Evaluation**: Models actively search for syntax errors, security issues, and performance problems
            - **Comprehensive Validation**: 7-metric assessment including syntax, security, and performance
            - **Edge Case Detection**: Identifies governor limits, bulkification issues, and B2B Commerce pitfalls
            - **Test Case Generation**: Automatic test class templates for migrated code
            - **Enhanced Error Detection**: Pattern-based syntax validation before AI processing
            
            **🤖 Model Roles:**
            - **Primary Model**: Performs initial conversion with skeptical analysis
            - **Validation Model**: Double-checks work with harsh but fair evaluation
            
            **⚠️ Important**: Always review and test AI-generated code in a sandbox before production deployment.
            """
        )

    app.launch()
    #return app




if __name__ == "__main__":
   main()


# print("✅ Initializing Salesforce Migration Assistant")
# app = main()
# print("✅ Application instance created")

# if __name__ == "__main__":
#     port = int(os.environ.get("PORT", 8080))
#     print(f"🚀 Starting server on port {port}")
#     try:
#         app.launch(
#             server_name="0.0.0.0",
#             server_port=port,
#             share=False
#         )
#     except Exception as e:
#         print(f"🔥 Server failed to start: {str(e)}")
#         raise
#     print("✅ Server started successfully")