from typing import Optional, Tuple
from os import getenv

import gradio as gr

from pydantic import BaseModel, Field
from langchain_openai import ChatOpenAI

def get_openai_key(api_key: Optional[str]) -> str:
    """Use the supplied API key or fallback to the environment variable."""
    if not api_key:
        api_key = getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OpenAI API key is required.")
    return api_key

# Define the schema for the Terraform configurations
class Terraform(BaseModel):
    """
    A Terraform module with a list of resources, provider, and configuration files.
    """
    main_tf: str = Field(..., description="The main Terraform configuration file.")
    variables_tf: str = Field(..., description="The Terraform variables configuration file.")
    outputs_tf: str = Field(..., description="The Terraform outputs configuration file.")

def generate(resources: str, provider: str, model: str, api_key: str = "") -> Tuple[str, str, str]:
    try:
        llm = ChatOpenAI(model=model, temperature=0.7, max_tokens=8192)
        structured_llm = llm.with_structured_output(Terraform)
        response = structured_llm.invoke(f"""You are an expert SRE assistant specialized in Terraform.

                       Generate a Terraform configuration for the following resources:
                       {resources} using the {provider} provider.

                       You always remember all of the necessary permissions and roles to configure for connecting the resources.
                       """)
        return response.main_tf, response.variables_tf, response.outputs_tf
    except Exception as e:
        raise ValueError(f"Failed to generate Terraform configuration: {e}")

# Define the Gradio interface
def create_interface():
    """Interface for the Terraform Module Generator."""
    with gr.Blocks() as demo:
        # Title
        gr.Markdown("## 🚀 Terraform Module Generator")

        # Instructional Markdown Box
        gr.Markdown("""
        **Welcome to the Terraform Module Generator!**

        This app assists you in generating Terraform configuration files based on your cloud provider and a natural language description of how your module should work.

        **How to Use:**
        1. **🔑 OpenAI API Key**: Enter your OpenAI API key to authenticate.
        2. **📦 Module Description**: Describe the key resources and how they work together.
        3. **🌐 Provider**: Select your cloud provider from the dropdown.
        4. **🤖 Language Model**: Choose the OpenAI model to use for generation.
        5. Click the **Generate Terraform Configuration** button to create your `.tf` files.

        The generated configurations will be displayed in separate tabs for easy access. There's a copy button in the upper right.
        """)
        with gr.Row():
            with gr.Column(scale=2):
                api_key = gr.Textbox(
                    label="🔑  OpenAI API Key",
                    type="password",
                    placeholder="Enter your OpenAI API key",
                )
            with gr.Column(scale=1):
                model = gr.Dropdown(
                    label="🤖  Language Model",
                    choices=["gpt-4o", "gpt-4o-mini"],
                    value="gpt-4o-mini",  # Default selection
                    type="value",
                )

        with gr.Row():
            resources = gr.Textbox(
                label="📦  Module Description",
                placeholder="Describe the resources and how they work together.",
                lines=5,
                scale=2,
            )
            provider = gr.Dropdown(
                label="☁ Cloud Provider",
                choices=["aws", "azure", "google"],
                value="aws",
                type="value",
                scale=1,
            )

        with gr.Row():
            generate_btn = gr.Button("⚙️ Generate Terraform Configuration")

        with gr.Tab("📄  main.tf"):
            main_tf_output = gr.Textbox(
                label="📄  main.tf",
                lines=20,
                show_copy_button=True,
                interactive=False
            )

        with gr.Tab("🔧  variables.tf"):
            variables_tf_output = gr.Textbox(
                label="🔧  variables.tf",
                lines=20,
                show_copy_button=True,
                interactive=False
            )

        with gr.Tab("📤  outputs.tf"):
            outputs_tf_output = gr.Textbox(
                label="📤  outputs.tf",
                lines=20,
                show_copy_button=True,
                interactive=False
            )

        # Define the button click event
        generate_btn.click(
            fn=generate,
            inputs=[resources, provider, model, api_key],
            outputs=[main_tf_output, variables_tf_output, outputs_tf_output]
        )

    return demo

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