import os
import base64
import io
from PIL import Image
from flask import Flask, request, jsonify
from flask_cors import CORS
from gradio_client import Client, handle_file
import openai

app = Flask(__name__)
CORS(app)

# Initialize CLIP Interrogator client
clipi_client = Client("https://fffiloni-clip-interrogator-2.hf.space/")

# Initialize LLM7 client
client = openai.OpenAI(
    base_url="https://api.llm7.io/v1",
    api_key=os.environ.get("LLM7_API_KEY", "unused")  # Use a free key or environment variable
)

def get_image_description(image_path):
    """Get image description using CLIP Interrogator"""
    try:
        print("Calling CLIP Interrogator...")
        result = clipi_client.predict(
            image=handle_file(image_path),
            mode="best",
            best_max_flavors=4,
            api_name="/clipi2"
        )
        print(f"CLIP description: {result}")
        return result
    except Exception as e:
        print(f"Error in get_image_description: {e}")
        return "a simple drawing"

def get_first_description(description):
    """Get only the first item from a comma-separated CLIP description"""
    items = [item.strip() for item in description.split(",")]
    return items[0] if items else description.strip()

def generate_story(description, audience="Children"):
    """Generate a kid-friendly story using GPT-5-Chat on LLM7 API"""
    first_desc = get_first_description(description)
    prompt = (
        f"Create a short, kid-friendly story for {audience} about: {first_desc}. "
        f"Use simple, cheerful words suitable for children. Include characters, action, "
        f"and make it imaginative."
        f"Write only 3 paragraphs."
        f"Do NOT add extra questions, suggestions, or prompts at the end."
    )

    print("Generating story with GPT-5-Chat...")
    try:
        response = client.chat.completions.create(
            model="gpt-5-chat",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.8
        )
        story = response.choices[0].message.content
        return story
    except Exception as e:
        print(f"Error generating story: {e}")
        return "Sorry, the story could not be generated."

@app.route("/health", methods=["GET"])
def health_check():
    return jsonify({"status": "healthy", "message": "Image-to-Story API is running"})

@app.route("/generate-story-base64", methods=["POST"])
def generate_story_base64():
    try:
        data = request.get_json()
        if "image" not in data:
            return jsonify({"error": "No image provided"}), 400

        # Decode base64
        try:
            image_data = base64.b64decode(data["image"])
            image = Image.open(io.BytesIO(image_data))
        except Exception:
            return jsonify({"error": "Invalid image data"}), 400

        if image.mode != "RGB":
            image = image.convert("RGB")

        temp_path = "temp_drawing.jpg"
        image.save(temp_path, "JPEG")

        audience = data.get("audience", "Children")

        # Get description + generate story
        description = get_image_description(temp_path)
        story = generate_story(description, audience)

        os.remove(temp_path)

        return jsonify({
            "success": True,
            "description": description,
            "story": story,
            "audience": audience
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

def create_app():
    return app

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port, debug=False)