import os
import gradio as gr
from langchain.chains import LLMChain, SequentialChain
from langchain.prompts import PromptTemplate
from langchain_community.llms import OpenAI  # <-- UPDATED IMPORT

# Get OpenAI API key from environment variable
openai_api_key = os.getenv("OPENAI_API_KEY")
llm = OpenAI(openai_api_key=openai_api_key, temperature=0.9)

# Define prompt templates
prompt_template_name = PromptTemplate(
    input_variables=['cuisine'],
    template="I want to open a restaurant for {cuisine} food. Suggest a fancy name for this"
)
name_chain = LLMChain(llm=llm, prompt=prompt_template_name, output_key="restaraunt_name")

prompt_template_items = PromptTemplate(
    input_variables=['restaraunt_name', 'cuisine'],
    template="Suggest a comprehensive menu for {restaraunt_name} which is {cuisine} cuisine. For each item, provide a brief, one-sentence description. Do not use any numbering or bullet points. Only mention items names and descriptions, break down according to category. Format the output as 'Category:\nItem Name - Description\n...'"
)
food_items_chain = LLMChain(llm=llm, prompt=prompt_template_items, output_key="menu_items")

chain = SequentialChain(
    chains=[name_chain, food_items_chain],
    input_variables=['cuisine'],
    output_variables=['restaraunt_name', 'menu_items']
)

TOP_CUISINES = [
    "Italian", "Chinese", "Japanese", "Mexican", "Indian",
    "French", "Thai", "Mediterranean", "American", "Spanish"
]

def generate_restaurant(cuisine):
    outputs = chain({'cuisine': cuisine})
    restaurant_name = outputs['restaraunt_name']
    menu_items_str = outputs['menu_items']

    lines = menu_items_str.strip().split('\n')
    menu_html = """
    <style>
        .menu-container {
            display: grid;
            grid-template-columns: repeat(3, 1fr);
            gap: 16px;
            padding: 16px;
            background-color: #f0e6d2;
            border-radius: 12px;
            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
        }
        .menu-tile {
            background-color: #fffaf0;
            border-radius: 8px;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
            padding: 16px;
            text-align: center;
            display: flex;
            flex-direction: column;
            justify-content: center;
            align-items: center;
            border: 1px solid #d2b48c;
        }
        .menu-tile h4 {
            margin: 0 0 8px 0;
            font-weight: bold;
            color: #5c4033;
            font-family: 'Comic Neue', cursive, sans-serif;
        }
        .menu-description {
            margin: 0;
            color: #8b4513;
            font-size: 0.9em;
            line-height: 1.4;
            text-align: left;
            font-family: 'Comic Neue', cursive, sans-serif;
        }
        .category-header {
            grid-column: 1 / -1;
            font-size: 1.5em;
            font-weight: bold;
            color: #a0522d;
            margin-top: 20px;
            margin-bottom: 10px;
            text-align: left;
            font-family: 'Comic Neue', cursive, sans-serif;
        }
    </style>
    <div class="menu-container">
    """
    current_category = ""
    for line in lines:
        line = line.strip()
        if not line:
            continue
        if ":" in line:
            current_category = line.replace(":", "")
            menu_html += f'<div class="category-header">{current_category}</div>'
        else:
            parts = line.split(' - ', 1)
            if len(parts) == 2:
                item_name, item_description = parts
                menu_html += f"""
                <div class="menu-tile">
                    <h4>{item_name.strip()}</h4>
                    <p class="menu-description">{item_description.strip()}</p>
                </div>
                """
    menu_html += "</div>"
    return restaurant_name, menu_html

def on_generate(selected, custom):
    cuisine = custom.strip() if custom and custom.strip() else selected
    return generate_restaurant(cuisine)

with gr.Blocks(title="Restaurant Business Generator", theme="soft", css="""
    @import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@400;700&display=swap');
    body {
        font-family: 'Comic Neue', 'Comic Sans MS', cursive, sans-serif;
    }
    .gradio-container {
        background-color: #f5f5dc;
        color: #5c4033;
    }
    .gradio-container h1, .gradio-container h2, .gradio-container h3, .gradio-container h4, .gradio-container h5, .gradio-container h6 {
        color: #5c4033;
    }
    .gradio-container label {
        color: #8b4513 !important;
        font-family: 'Comic Neue', cursive, sans-serif;
    }
    .gradio-container .gr-markdown p {
        color: #8b4513 !important;
    }
    .gradio-container input[type="text"], .gradio-container textarea {
        border-color: #d2b48c;
    }
    .gradio-container .gr-dropdown {
        border-color: #d2b48c;
    }
    .gradio-container .gr-dropdown-selected {
        background-color: #fffaf0;
    }
    #generate-btn {
        background-color: #a0522d;
        color: white;
        font-weight: bold;
        border: none;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
    }
    #generate-btn:hover {
        background-color: #8b4513;
    }
    .output-class {
        background-color: #faebd7;
        border: 1px solid #d2b48c;
        border-radius: 8px;
        padding: 10px;
        box-shadow: 0 4px 6px rgba(0, 0, 0, 0.08);
    }
    #restaurant-name-output textarea {
        color: #5c4033;
        font-weight: bold;
        font-size: 1.1em;
        background-color: #fffaf0 !important;
        border: 1px solid #d2b48c !important;
    }
    :root {
        --primary-50: #fffaf0 !important;
        --primary-100: #fff5e1 !important;
        --primary-200: #fff0d2 !important;
        --primary-300: #ffe9b9 !important;
        --primary-400: #ffe1a0 !important;
        --primary-500: #a0522d !important;
        --primary-600: #8b4513 !important;
        --primary-700: #5c4033 !important;
        --primary-800: #4a332a !important;
        --primary-900: #3b2820 !important;
        --primary-950: #291a14 !important;
    }
""") as demo:
    gr.Markdown("<h1 style='text-align:center; color: #5c4033;'>✨ The Magical Restaurant Idea Generator ✨</h1>")
    with gr.Row():
        with gr.Column(scale=1, min_width=200):
            gr.Markdown("<p style='color:#8b4513; text-align:center;'>Choose a cuisine or type your own, and let the magic happen! We'll give you a cute restaurant name and a delightful menu.</p>")
            cuisine_dropdown = gr.Dropdown(choices=TOP_CUISINES, label="Popular Cuisines", value=TOP_CUISINES[0])
            cuisine_textbox = gr.Textbox(label="Or enter a custom cuisine", placeholder="e.g. Vietnamese, Peruvian")
            generate_btn = gr.Button("Create My Restaurant! 🪄", variant="primary", elem_id="generate-btn")
        with gr.Column(scale=2, elem_classes="output-class"):
            name_output = gr.Textbox(label="Suggested Restaurant Name", interactive=False, elem_id="restaurant-name-output")
            menu_output = gr.HTML(label="Suggested Menu Items")

    generate_btn.click(fn=on_generate, inputs=[cuisine_dropdown, cuisine_textbox], outputs=[name_output, menu_output])

if __name__ == "__main__":
    demo.launch()