import gradio as gr
import requests
import os
from typing import Dict, Tuple

# Constants
DEFAULTS = {
    "SYSTEM_MESSAGE": "You are CodePal.ai, an expert AI assistant specialized in helping programmers. You generate clean, efficient, and well-documented code based on user requirements.",
    "MAX_TOKENS": 4000,
    "TEMPERATURE": 0.7,
    "TOP_P": 0.95,
}

API_URLS = {
    "CODE_GENERATOR": "https://api.codepal.ai/v1/code-generator/query",
    "CODE_FIXER": "https://api.codepal.ai/v1/code-fixer/query",
    "CODE_EXTENDER": "https://api.codepal.ai/v1/code-extender/query",
}

def get_api_key() -> str:
    """Fetch API key from environment variables."""
    return os.environ.get("CODEPAL_API_KEY", "")

def is_api_key_valid() -> bool:
    """Validate if the API key is set and not a placeholder."""
    api_key = get_api_key()
    return bool(api_key) and api_key != "YOUR_API_KEY"

def build_request_data(language: str, instructions: str, flavor: str, max_tokens: int, temperature: float, top_p: float) -> Dict:
    """Construct the request data for API calls."""
    return {
        "language": language,
        "instructions": instructions,
        "flavor": flavor,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "top_p": top_p
    }

def handle_api_response(response) -> Tuple[bool, str]:
    """Handle API response and return success status and result."""
    if response.status_code != 200:
        error_message = "API request failed"
        try:
            error_data = response.json()
            error_message = error_data.get("error", error_message)
        except Exception:
            error_message = f"API request failed with status code {response.status_code}"
        return False, f"Error: {error_message}"

    result = response.json()
    if "error" in result:
        return False, f"Error: {result['error']}"
    
    return True, result["result"]

def generate_code(language: str, requirements: str, code_style: str, include_tests: bool, 
                  max_tokens: int = DEFAULTS["MAX_TOKENS"], 
                  temperature: float = DEFAULTS["TEMPERATURE"], 
                  top_p: float = DEFAULTS["TOP_P"]) -> Tuple[bool, str]:
    """Generate code using CodePal.ai API."""
    if not is_api_key_valid():
        return False, "Error: CodePal.ai API key is not configured. Please set the CODEPAL_API_KEY environment variable."
    
    flavor = {
        "minimal": "minimal",
        "verbose": "documented"
    }.get(code_style, "standard") if code_style in ["minimal", "verbose"] else "tests" if include_tests else "standard"

    api_key = get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    data = build_request_data(language, requirements, flavor, max_tokens, temperature, top_p)

    try:
        response = requests.post(API_URLS["CODE_GENERATOR"], headers=headers, json=data)
        return handle_api_response(response)
    except Exception as e:
        return False, f"Error: {str(e)}"

def fix_code(code: str, instructions: str, language: str = "", 
             max_tokens: int = DEFAULTS["MAX_TOKENS"], 
             temperature: float = DEFAULTS["TEMPERATURE"], 
             top_p: float = DEFAULTS["TOP_P"]) -> Tuple[bool, str]:
    """Fix existing code using CodePal.ai API."""
    if not is_api_key_valid():
        return False, "Error: CodePal.ai API key is not configured. Please set the CODEPAL_API_KEY environment variable."
    
    api_key = get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    data = build_request_data(language, instructions, "", max_tokens, temperature, top_p)
    if language and language != "auto-detect":
        data["language"] = language

    try:
        response = requests.post(API_URLS["CODE_FIXER"], headers=headers, json=data)
        return handle_api_response(response)
    except Exception as e:
        return False, f"Error: {str(e)}"

def extend_code(code: str, instructions: str, language: str = "", 
                max_tokens: int = DEFAULTS["MAX_TOKENS"], 
                temperature: float = DEFAULTS["TEMPERATURE"], 
                top_p: float = DEFAULTS["TOP_P"]) -> Tuple[bool, str]:
    """Extend existing code using CodePal.ai API."""
    if not is_api_key_valid():
        return False, "Error: CodePal.ai API key is not configured. Please set the CODEPAL_API_KEY environment variable."
    
    api_key = get_api_key()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    
    data = build_request_data(language, instructions, "", max_tokens, temperature, top_p)
    if language and language != "auto-detect":
        data["language"] = language

    try:
        response = requests.post(API_URLS["CODE_EXTENDER"], headers=headers, json=data)
        return handle_api_response(response)
    except Exception as e:
        return False, f"Error: {str(e)}"

def save_to_file(code: str, language: str, file_name: str = "generated_code") -> str:
    """Save the generated code to a file with appropriate extension."""
    extensions = {
        "python": ".py", "javascript": ".js", "typescript": ".ts", "java": ".java",
        "c": ".c", "cpp": ".cpp", "csharp": ".cs", "go": ".go", "rust": ".rs",
        "php": ".php", "ruby": ".rb", "swift": ".swift", "kotlin": ".kt",
        "html": ".html", "css": ".css", "sql": ".sql", "shell": ".sh",
    }
    
    extension = extensions.get(language.lower(), ".txt")
    full_file_name = f"{file_name}{extension}"
    
    try:
        with open(full_file_name, "w") as f:
            f.write(code)
        return f"Code successfully saved to {full_file_name}"
    except Exception as e:
        return f"Error saving code to file: {str(e)}"

def ui_generate_code(language: str, generation_type: str, requirements: str, code_style: str, 
                     dependencies: str, include_tests: bool, 
                     max_tokens: int = DEFAULTS["MAX_TOKENS"],
                     temperature: float = DEFAULTS["TEMPERATURE"],
                     top_p: float = DEFAULTS["TOP_P"],
                     save_file: bool = False) -> str:
    """UI handler for code generation."""
    combined_requirements = f"{generation_type}: {requirements}"
    if dependencies:
        combined_requirements += f"\nDependencies: {dependencies}"
    
    success, result = generate_code(language, combined_requirements, code_style, include_tests, 
                                     max_tokens, temperature, top_p)
    
    if not success:
        return result
    
    if save_file:
        file_message = save_to_file(result, language)
        result = f"{result}\n\n{file_message}"
    
    return result

def ui_fix_code(code: str, instructions: str, language: str, 
                 max_tokens: int = DEFAULTS["MAX_TOKENS"],
                 temperature: float = DEFAULTS["TEMPERATURE"],
                 top_p: float = DEFAULTS["TOP_P"],
                 save_file: bool = False) -> str:
    """UI handler for code fixing."""
    success, result = fix_code(code, instructions, language, max_tokens, temperature, top_p)
    
    if not success:
        return result
    
    if save_file:
        file_message = save_to_file(result, language, "fixed_code")
        result = f"{result}\n\n{file_message}"
    
    return result

def ui_extend_code(code: str, instructions: str, language: str, 
                   max_tokens: int = DEFAULTS["MAX_TOKENS"],
                   temperature: float = DEFAULTS["TEMPERATURE"],
                   top_p: float = DEFAULTS["TOP_P"],
                   save_file: bool = False) -> str:
    """UI handler for code extension."""
    success, result = extend_code(code, instructions, language, max_tokens, temperature, top_p)
    
    if not success:
        return result
    
    if save_file:
        file_message = save_to_file(result, language, "extended_code")
        result = f"{result}\n\n{file_message}"
    
    return result

# Create Gradio Interface with tabs for different functionality
with gr.Blocks(title="CodePal.ai Interface") as app:
    gr.Markdown(
        """
        # CodePal.ai Code Generator
        
        A custom interface for the CodePal.ai API providing code generation, fixing, and extension capabilities.
        
        **Note**: You need to set the `CODEPAL_API_KEY` environment variable with your CodePal.ai API key.
        """
    )
    
    api_key_status = gr.Markdown(
        "⚠️ **API Key Status**: Checking..." if is_api_key_valid() else "❌ **API Key Status**: Missing. Please set the CODEPAL_API_KEY environment variable."
    )
    
    with gr.Tabs():
        # Generate Code Tab
        with gr.TabItem("Generate Code"):
            with gr.Row():
                with gr.Column(scale=1):
                    language = gr.Dropdown(
                        choices=["javascript", "typescript", "python", "react", "vue", "node", "php", "java", "csharp", "html", "css", "ruby", "go", "rust", "swift"],
                        value="javascript",
                        label="Language/Framework"
                    )
                    
                    generation_type = gr.Radio(
                        choices=["component", "function", "api", "algorithm", "style", "other"],
                        value="function",
                        label="Generation Type"
                    )
                    
                    requirements = gr.Textbox(
                        lines=5,
                        placeholder="Describe what you need in detail. Example: Create a function that sorts an array of objects by a specific property.",
                        label="Requirements",
                        info="Be as specific as possible for better results."
                    )
                    
                    with gr.Accordion("Advanced Options", open=False):
                        code_style = gr.Dropdown(
                            choices=["standard", "functional", "oop", "minimal", "verbose"],
                            value="standard",
                            label="Code Style"
                        )
                        
                        dependencies = gr.Textbox(
                            lines=2,
                            placeholder="E.g., react-router-dom, lodash, axios",
                            label="Dependencies (optional)",
                            info="List any packages or libraries you want to use."
                        )
                        
                        include_tests = gr.Checkbox(
                            label="Include unit tests",
                            value=False
                        )
                        
                        with gr.Group():
                            gr.Markdown("### Model Settings")
                            max_tokens = gr.Slider(
                                minimum=1000, 
                                maximum=8000, 
                                step=100, 
                                value=DEFAULTS["MAX_TOKENS"],
                                label="Max Tokens",
                                info="Maximum number of tokens (words) in the generated response"
                            )
                            
                            temperature = gr.Slider(
                                minimum=0.1, 
                                maximum=1.0, 
                                step=0.05, 
                                value=DEFAULTS["TEMPERATURE"],
                                label="Temperature",
                                info="Higher values make output more random, lower values make it more deterministic"
                            )
                            
                            top_p = gr.Slider(
                                minimum=0.5, 
                                maximum=1.0, 
                                step=0.05, 
                                value=DEFAULTS["TOP_P"],
                                label="Top P",
                                info="Controls diversity via nucleus sampling"
                            )
                            
                            save_gen_file = gr.Checkbox(
                                label="Save generated code to file",
                                value=False,
                                info="Automatically save the generated code to a file with appropriate extension"
                            )
                    
                    generate_btn = gr.Button("Generate Code", variant="primary")
                
                with gr.Column(scale=1):
                    output_code = gr.Code(
                        label="Generated Code",
                        language="javascript",
                        interactive=False,
                        lines=20
                    )
            
            generate_btn.click(
                fn=ui_generate_code,
                inputs=[language, generation_type, requirements, code_style, dependencies, include_tests, 
                        max_tokens, temperature, top_p],
                outputs=[output_code]
            )
            
            language.change(
                fn=lambda lang: gr.update(language=lang.lower()),
                inputs=[language],
                outputs=[output_code]
            )
        
        # Fix Code Tab
        with gr.TabItem("Fix Code"):
            with gr.Row():
                with gr.Column(scale=1):
                    fix_language = gr.Dropdown(
                        choices=["auto-detect", "javascript", "typescript", "python", "react", "vue", "node", "php", "java", "csharp", "html", "css", "ruby", "go", "rust", "swift"],
                        value="auto-detect",
                        label="Language (optional)"
                    )
                    
                    original_code = gr.Code(
                        label="Code to Fix",
                        language="javascript",
                        lines=10
                    )
                    
                    fix_instructions = gr.Textbox(
                        lines=3,
                        placeholder="Describe what's wrong or what needs to be fixed.",
                        label="Fix Instructions",
                        info="Be specific about the issues or errors you want resolved."
                    )
                    
                    with gr.Accordion("Advanced Model Settings", open=False):
                        fix_max_tokens = gr.Slider(
                            minimum=1000, 
                            maximum=8000, 
                            step=100, 
                            value=DEFAULTS["MAX_TOKENS"],
                            label="Max Tokens",
                            info="Maximum number of tokens (words) in the generated response"
                        )
                        
                        fix_temperature = gr.Slider(
                            minimum=0.1, 
                            maximum=1.0, 
                            step=0.05, 
                            value=DEFAULTS["TEMPERATURE"],
                            label="Temperature",
                            info="Higher values make output more random, lower values make it more deterministic"
                        )
                        
                        fix_top_p = gr.Slider(
                            minimum=0.5, 
                            maximum=1.0, 
                            step=0.05, 
                            value=DEFAULTS["TOP_P"],
                            label="Top P",
                            info="Controls diversity via nucleus sampling"
                        )
                        
                        save_fix_file = gr.Checkbox(
                            label="Save fixed code to file",
                            value=False,
                            info="Automatically save the fixed code to a file with appropriate extension"
                        )
                    
                    fix_btn = gr.Button("Fix Code", variant="primary")
                
                with gr.Column(scale=1):
                    fixed_code = gr.Code(
                        label="Fixed Code",
                        language="javascript",
                        interactive=False,
                        lines=20
                    )
            
            fix_btn.click(
                fn=ui_fix_code,
                inputs=[original_code, fix_instructions, fix_language,
                        fix_max_tokens, fix_temperature, fix_top_p],
                outputs=[fixed_code]
            )
            
            fix_language.change(
                fn=lambda lang: gr.update(language=lang.lower() if lang != "auto-detect" else "javascript"),
                inputs=[fix_language],
                outputs=[fixed_code]
            )
        
        # Extend Code Tab
        with gr.TabItem("Extend Code"):
            with gr.Row():
                with gr.Column(scale=1):
                    extend_language = gr.Dropdown(
                        choices=["auto-detect", "javascript", "typescript", "python", "react", "vue", "node", "php", "java", "csharp", "html", "css", "ruby", "go", "rust", "swift"],
                        value="auto-detect",
                        label="Language (optional)"
                    )
                    
                    base_code = gr.Code(
                        label="Base Code",
                        language="javascript",
                        lines=10
                    )
                    
                    extend_instructions = gr.Textbox(
                        lines=3,
                        placeholder="Describe how you want to extend the code.",
                        label="Extension Instructions",
                        info="Be specific about what new functionality or features you want to add."
                    )
                    
                    with gr.Accordion("Advanced Model Settings", open=False):
                        extend_max_tokens = gr.Slider(
                            minimum=1000, 
                            maximum=8000, 
                            step=100, 
                            value=DEFAULTS["MAX_TOKENS"],
                            label="Max Tokens",
                            info="Maximum number of tokens (words) in the generated response"
                        )
                        
                        extend_temperature = gr.Slider(
                            minimum=0.1, 
                            maximum=1.0, 
                            step=0.05, 
                            value=DEFAULTS["TEMPERATURE"],
                            label="Temperature",
                            info="Higher values make output more random, lower values make it more deterministic"
                        )
                        
                        extend_top_p = gr.Slider(
                            minimum=0.5, 
                            maximum=1.0, 
                            step=0.05, 
                            value=DEFAULTS["TOP_P"],
                            label="Top P",
                            info="Controls diversity via nucleus sampling"
                        )
                    
                    extend_btn = gr.Button("Extend Code", variant="primary")
                
                with gr.Column(scale=1):
                    extended_code = gr.Code(
                        label="Extended Code",
                        language="javascript",
                        interactive=False,
                        lines=20
                    )
            
            extend_btn.click(
                fn=ui_extend_code,
                inputs=[base_code, extend_instructions, extend_language,
                        extend_max_tokens, extend_temperature, extend_top_p],
                outputs=[extended_code]
            )
            
            extend_language.change(
                fn=lambda lang: gr.update(language=lang.lower() if lang != "auto-detect" else "javascript"),
                inputs=[extend_language],
                outputs=[extended_code]
            )

# Launch the app
if __name__ == "__main__":
    import sys
    port = int(os.environ.get("GRADIO_SERVER_PORT", 5000))
    alt_port = 6000
    
    try:
        app.launch(
            server_name="0.0.0.0",
            server_port=port,
            share=True,
            debug=True
        )
    except OSError:
        print(f"Port {port} is in use, trying port {alt_port} instead...")
        app.launch(
            server_name="0.0.0.0",
            server_port=alt_port,
            share=True,
            debug=True
        )
    except Exception as e:
        print(f"Error starting the Gradio application: {str(e)}")
        sys.exit(1)