import gradio as gr
from transformers import pipeline
import random

# Initialize the Hugging Face text generation pipeline with distilgpt2
generator = pipeline("text-generation", model="distilgpt2")

# Function to generate checklists, tips, and engagement score
def generate_project_data(project_input):
    # Generate checklists (3 tasks)
    checklist_prompt = f"Generate a list of 3 safety and productivity tasks for a construction project: {project_input}"
    checklist_response = generator(checklist_prompt, max_length=100, num_return_sequences=1, truncation=True)[0]["generated_text"]
    # Extract tasks (simple parsing assuming the model returns a list-like structure)
    tasks = checklist_response.replace(checklist_prompt, "").split(".")[:3]
    tasks = [task.strip() for task in tasks if task.strip()]
    if len(tasks) < 3:
        # Fallback tasks if the model doesn't generate enough
        tasks.extend([
            "Conduct a safety briefing with the team.",
            "Inspect all equipment before use.",
            "Ensure all workers are wearing PPE."
        ][:3 - len(tasks)])

    # Generate a tip
    tip_prompt = f"Provide a productivity tip for a construction project supervisor: {project_input}"
    tip_response = generator(tip_prompt, max_length=50, num_return_sequences=1, truncation=True)[0]["generated_text"]
    tip = tip_response.replace(tip_prompt, "").strip()
    if not tip:
        tip = "Schedule regular breaks to maintain team focus."

    # Generate a mock engagement score (rule-based for simplicity)
    # In a real scenario, this could be generated by a model trained on engagement data
    engagement_score = random.randint(70, 90)  # Random score between 70 and 90

    # Return the data in the expected JSON format
    return {
        "checklists": [{"task": task} for task in tasks],
        "tips": tip,
        "engagementScore": engagement_score
    }

# Create a Gradio interface
interface = gr.Interface(
    fn=generate_project_data,
    inputs=gr.Textbox(label="Project Input", placeholder="Enter project details (e.g., Project: Highway Construction, Start Date: 2025-05-01)"),
    outputs=gr.JSON(label="Generated Data"),
    title="AI Coach Data Generator",
    description="Generates daily checklists, tips, and engagement scores for construction projects."
)

# Launch the app
interface.launch()