import os
import gradio as gr
from crewai import Agent, Task, Crew, Process
# CORRECTED IMPORTS FOR NEWER VERSIONS
from crewai_tools import DuckDuckGoSearchRun, FileWriterTool
# ---
from langchain_google_genai import ChatGoogleGenerativeAI

# --- 1. SET UP THE GEMINI LLM ---
try:
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash",
                                 verbose=True,
                                 temperature=0.6,
                                 google_api_key=os.environ.get("GOOGLE_API_KEY"))
except Exception as e:
    print("\n\n---")
    print("Error initializing Gemini LLM. Make sure you have set your GOOGLE_API_KEY.")
    print(f"Error: {e}")
    print("---\n\n")
    llm = None

# --- 2. DEFINE TOOLS AND OUTPUT DIRECTORY ---
search_tool = DuckDuckGoSearchRun()
os.makedirs("outputs", exist_ok=True)
file_write_tool = FileWriterTool(file_path="outputs/index.html")

# --- 3. DEFINE THE AGENT TEAM ---
designer = Agent(
  role='Principal Brand Designer',
  goal='Translate a user\'s "vibe" into a concrete, professional design brief for a website.',
  backstory=(
      "You are a world-class brand designer with a keen eye for aesthetics. "
      "You excel at understanding abstract concepts and turning them into actionable design elements."
  ),
  verbose=True,
  llm=llm,
  tools=[search_tool],
  allow_delegation=False
)
copywriter = Agent(
  role='Senior UX Copywriter',
  goal='Write compelling, vibe-appropriate copy for all sections of a website based on a design brief.',
  backstory=(
      "You are an expert copywriter who specializes in creating website content that resonates with the target audience."
  ),
  verbose=True,
  llm=llm,
  allow_delegation=False
)
developer = Agent(
  role='Senior Frontend Developer',
  goal='Take a design brief and content map to build a complete, self-contained HTML file.',
  backstory=(
      "You are a skilled frontend developer who can build beautiful and functional websites from scratch."
  ),
  verbose=True,
  llm=llm,
  tools=[file_write_tool],
  allow_delegation=False
)

# --- 4. DEFINE THE WORKFLOW (TASKS) ---
design_task = Task(
  description=(
      "Analyze the user's prompt: '{prompt}' to understand the desired vibe. "
      "Conduct a web search for inspiration on color palettes, font pairings, and overall aesthetics. "
      "Create a comprehensive 'Design Brief' document with a color palette (hex codes), font pairing (Google Fonts), and a detailed page structure."
  ),
  expected_output=(
      "A structured markdown document with the complete Design Brief (Color Palette, Font Pairing, Page Structure)."
  ),
  agent=designer
)
copywriting_task = Task(
  description=(
      "Using the Design Brief, write all text content for the website. "
      "Ensure the tone matches the vibe. Create placeholder content for all sections."
  ),
  expected_output=(
      "A structured markdown 'Content Map' with text for each section of the website."
  ),
  agent=copywriter,
  context=[design_task]
)
development_task = Task(
  description=(
      "Based on the Design Brief and Content Map, create a single-file HTML website. "
      "The file must be self-contained with embedded CSS and any simple JavaScript. "
      "Use the specified colors and Google Fonts. Your final output MUST be only the code for the HTML file."
  ),
  expected_output=(
      "The complete HTML code for the website, written to 'outputs/index.html'."
  ),
  agent=developer,
  context=[design_task, copywriting_task],
  tools=[file_write_tool]
)

# --- 5. ASSEMBLE THE CREW ---
vibe_crew = Crew(
  agents=[designer, copywriter, developer],
  tasks=[design_task, copywriting_task, development_task],
  process=Process.sequential,
  verbose=2
)

# --- 6. CREATE THE GRADIO WEB INTERFACE ---
def run_crew(prompt):
    if llm is None:
        return "Error: Gemini LLM not initialized. Please check your API key.", None, None
    if os.path.exists("outputs/index.html"):
        os.remove("outputs/index.html")
    inputs = {'prompt': prompt}
    result = vibe_crew.kickoff(inputs=inputs)
    try:
        with open("outputs/index.html", "r") as file:
            html_content = file.read()
        return result, html_content, "outputs/index.html"
    except FileNotFoundError:
        return result, "Error: Output file not found. The developer agent might have failed.", None

iface = gr.Interface(
    fn=run_crew,
    inputs=gr.Textbox(
        lines=3,
        placeholder="e.g., A website for my new coffee shop in Brooklyn. The vibe should be cozy, rustic, and artisanal."
    ),
    outputs=[
        gr.Markdown(label="Crew Final Report"),
        gr.Code(label="Generated HTML Code", language="html"),
        gr.File(label="Download Website")
    ],
    title="🎨 Vibe Coder AI",
    description="Enter the 'vibe' for your website, and this multi-agent team will design, write, and code a complete webpage for you using Gemini.",
    examples=[
        ["A personal portfolio for a photographer. The vibe should be minimalist, modern, and clean."],
        ["A landing page for a new tech startup focused on AI. Vibe: futuristic, sleek, professional."],
        ["A website for a local bakery. Vibe: warm, friendly, and homemade."]
    ],
    allow_flagging="never"
)

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