import os
import random
import pickle
import re
import tempfile
import numpy as np
import cv2
import gradio as gr
from PIL import Image
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForImageClassification, AutoImageProcessor
from gradio_client import Client, handle_file
import openai
from flask import Flask, request, jsonify, Response
from flask_cors import CORS
import base64
import io

# ============================================================
# INITIALIZATION
# ============================================================
print("Initializing Tree & House Story Generator (Flask + Gradio)")

# OpenAI
client = openai.OpenAI(
    base_url="https://api.llm7.io/v1",
    api_key=os.environ.get("LLM7_API_KEY", "unused")
)

# NSFW MODEL
nsfw_processor = AutoImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")
nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")

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

# SAFETY
UNSAFE_KEYWORDS = [
    "penis", "vagina", "genitals", "testicle",
    "sex", "sexual", "erotic", "porn", "explicit",
    "nipple", "genital", "lewd"
]

def is_safe_description(text):
    text = text.lower()
    return not any(re.search(rf"\b{w}\b", text) for w in UNSAFE_KEYWORDS)

def is_drawing_safe(image):
    try:
        inputs = nsfw_processor(images=image, return_tensors="pt")
        outputs = nsfw_model(**inputs)
        probs = outputs.logits.softmax(dim=1)[0]
        if probs[1] > 0.85:
            return False, "NSFW detected"
        return True, None
    except:
        return True, None

# LOAD TREE & HOUSE MODEL
classifier = None
bert_model = None
tree_stories = []
house_stories = []

try:
    with open("tree_house_model.pkl", "rb") as f:
        model_package = pickle.load(f)

    classifier = model_package["classifier"]
    bert_model = SentenceTransformer(model_package["bert_model_name"])
    tree_stories = model_package["tree_stories"]
    house_stories = model_package["house_stories"]
    model_accuracy = model_package["accuracy"]

    print(f"✓ Model loaded ({model_accuracy*100:.2f}%)")
except:
    print("Tree/House model not loaded")

# CORE FUNCTIONS
def get_image_description(image):
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
        path = f.name
        image.save(path)
    try:
        desc = clipi_client.predict(
            image=handle_file(path),
            mode="best",
            api_name="/clipi2"
        )
    finally:
        os.remove(path)
    return desc

def predict_category(description):
    if not classifier:
        return None, 0.0, {}
    emb = bert_model.encode([description], normalize_embeddings=True)
    pred = classifier.predict(emb)[0]
    probs = classifier.predict_proba(emb)[0]
    return pred, float(max(probs)), {"tree": float(probs[0]), "house": float(probs[1])}

def get_random_story(category):
    stories = tree_stories if category == "tree" else house_stories
    if not stories:
        return None
    s = random.choice(stories).copy()
    s.pop("category", None)
    return s

def generate_story_openai(description, audience, is_vulgar=False):
    """
    description: user drawing description (string or list)
    audience: string, e.g., 'children aged 5-7'
    is_vulgar: boolean, True if Hugging Face classifier detects vulgar content
    """

    # Handle list input
    if isinstance(description, list):
        first_desc = description[0]
    else:
        first_desc = description

    # Pre-check using local unsafe keywords
    lower_desc = first_desc.lower()
    for word in UNSAFE_KEYWORDS:
        if word in lower_desc:
            return "Your drawing contains inappropriate content and cannot be used to generate a story."

    # If Hugging Face classifier detects vulgar content
    if is_vulgar:
        return "Your drawing contains inappropriate content and cannot be used to generate a story."

    # Prompt with explicit instruction
    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"Also, provide a creative title at the very beginning. "
        f"IMPORTANT: If the input contains any vulgar, sexual, or inappropriate content, "
        f"do NOT generate any story and reply only with: "
        f"'The content is inappropriate and cannot be used for a story.' "
        f"Do NOT add extra questions, suggestions, or prompts at the end."
    )

    try:
        res = client.chat.completions.create(
            model="gpt-5-chat",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.8
        )
        return res.choices[0].message.content
    except:
        return "Story generation failed."


def format_story(story_data):
    """Return story text ONLY – no labels like characters, setting, plot, etc."""

    if isinstance(story_data, str):
        return story_data

    elif isinstance(story_data, dict):
        parts = []

        # Optional title (pwede mo rin alisin kung ayaw mo)
        if "title" in story_data and story_data["title"]:
            parts.append(story_data["title"].strip())

        # All possible story parts (NO LABELS)
        ordered_keys = [
            "characters",
            "setting",
            "plot",
            "problem",
            "ending",
            "lesson",
            "moral",
            "story"
        ]

        for key in ordered_keys:
            if key in story_data and story_data[key]:
                parts.append(story_data[key].strip())

        # Join as clean paragraphs
        return "\n\n".join(parts)

    else:
        return str(story_data)

def process_image_base64(image_base64, audience="Children"):
    image_data = base64.b64decode(image_base64)
    image = Image.open(io.BytesIO(image_data)).convert("RGB")

    # Description mula sa CLIP
    description = get_image_description(image)
    
    # Predict category (tree, house, other)
    category, confidence, probs = predict_category(description)

    # Check if description contains tree/house keywords
    description_lower = description.lower()
    is_tree = "tree" in description_lower or "branch" in description_lower or "leaf" in description_lower or "trunk" in description_lower
    is_house = "house" in description_lower or "building" in description_lower or "home" in description_lower or "roof" in description_lower

    # TREE/HOUSE → always use dataset story, ignore vulgar/NSFW
    if (category == "tree" and is_tree) or (category == "house" and is_house):
        story_data = get_random_story("tree" if category == "tree" else "house")
        story_category = category
        formatted_story = format_story(story_data)
        return {
            "success": True,
            "description": description,
            "category": story_category,
            "confidence": confidence,
            "probabilities": probs,
            "story": formatted_story
        }

    # Ibang drawings → GPT-generated, check NSFW and vulgar keywords
    safe, reason = is_drawing_safe(image)
    if not safe:
        return {"success": False, "error": "Bawal sa bata: NSFW detected", "reason": reason}

    if not is_safe_description(description):
        return {"success": False, "error": "Bawal sa bata: unsafe description detected"}

    formatted_story = generate_story_openai(description, audience)
    story_category = "OPENAI"

    return {
        "success": True,
        "description": description,
        "category": story_category,
        "confidence": None,
        "probabilities": probs,
        "story": formatted_story
    }



# GRADIO FUNCTION
def generate_story_gradio(image, audience):
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_b64 = base64.b64encode(buffered.getvalue()).decode()
    result = process_image_base64(img_b64, audience)
    if not result["success"]:
        return result.get("error", "Failed"), "", "", ""
    # Story is already formatted as a string
    return result["description"], result["category"], str(result["probabilities"]), result["story"]

# GRADIO UI
with gr.Blocks(title=" Tree & House Story Generator") as demo:
    gr.Markdown("# Tree & House Story Generator")
    gr.Markdown("Upload a drawing and generate a creative story!")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Upload a drawing")
            audience_input = gr.Dropdown(["Children", "Teens"], value="Children", label="Audience")
            btn = gr.Button(" Generate Story", variant="primary")
        
        with gr.Column():
            description_output = gr.Textbox(label="Image Description", lines=2)
            category_output = gr.Textbox(label="Category")
            probabilities_output = gr.Textbox(label="Probabilities")
            story_output = gr.Textbox(label="Story", lines=12)
    
    btn.click(
        generate_story_gradio,
        inputs=[image_input, audience_input],
        outputs=[description_output, category_output, probabilities_output, story_output]
    )

# FLASK API
app = Flask(__name__)
CORS(app)

# Create Gradio WSGI app with SSR disabled
gradio_app = gr.routes.App.create_app(demo, ssr_mode=False)

@app.route('/')
def index():
    return Response("Redirecting to Gradio UI...", status=302, headers={"Location": "/gradio"})

@app.route("/api/health", methods=["GET"])
def health_check():
    return jsonify({
        "status": "healthy",
        "model_loaded": classifier is not None,
        "tree_stories": len(tree_stories),
        "house_stories": len(house_stories)
    })

@app.route("/api/generate-story-base64", methods=["POST"])
def generate_story_api():
    data = request.get_json()
    if "image" not in data:
        return jsonify({"error": "No image provided"}), 400
    audience = data.get("audience", "Children")
    result = process_image_base64(data["image"], audience)
    return jsonify(result)

@app.route('/api/test', methods=['GET'])
def api_test():
    """API documentation and test page"""
    html = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>API Test - Story Generator</title>
        <style>
            body { font-family: Arial, sans-serif; max-width: 800px; margin: 50px auto; padding: 20px; }
            h1 { color: #333; }
            .endpoint { background: #f5f5f5; padding: 15px; margin: 20px 0; border-radius: 5px; }
            .method { display: inline-block; padding: 5px 10px; border-radius: 3px; font-weight: bold; }
            .get { background: #61affe; color: white; }
            .post { background: #49cc90; color: white; }
            code { background: #f5f5f5; padding: 2px 5px; border-radius: 3px; }
            .test-section { margin: 20px 0; padding: 20px; border: 1px solid #ddd; border-radius: 5px; }
            button { background: #49cc90; color: white; border: none; padding: 10px 20px; cursor: pointer; border-radius: 5px; }
            button:hover { background: #3da877; }
            #result { margin-top: 15px; padding: 15px; background: #f9f9f9; border-radius: 5px; white-space: pre-wrap; }
        </style>
    </head>
    <body>
        <h1> Story Generator API</h1>
        
        <div class="endpoint">
            <span class="method get">GET</span> <code>/api/health</code>
            <p>Check API status</p>
        </div>
        
        <div class="endpoint">
            <span class="method post">POST</span> <code>/api/generate-story-base64</code>
            <p>Generate story from base64 image</p>
            <strong>Request body:</strong>
            <pre>{
  "image": "base64_encoded_image_string",
  "audience": "Children"
}</pre>
        </div>

        <div class="test-section">
            <h2>Test API</h2>
            <p>Upload an image to test the story generation:</p>
            <input type="file" id="imageInput" accept="image/*">
            <br><br>
            <label>Audience: 
                <select id="audience">
                    <option>Children</option>
                    <option>Teens</option>
                </select>
            </label>
            <br><br>
            <button onclick="testAPI()">Generate Story</button>
            <div id="result"></div>
        </div>

        <script>
            async function testAPI() {
                const fileInput = document.getElementById('imageInput');
                const audience = document.getElementById('audience').value;
                const resultDiv = document.getElementById('result');
                
                if (!fileInput.files[0]) {
                    resultDiv.textContent = 'Please select an image first!';
                    return;
                }
                
                resultDiv.textContent = 'Processing...';
                
                const reader = new FileReader();
                reader.onload = async function(e) {
                    const base64 = e.target.result.split(',')[1];
                    
                    try {
                        const response = await fetch('/api/generate-story-base64', {
                            method: 'POST',
                            headers: {
                                'Content-Type': 'application/json'
                            },
                            body: JSON.stringify({
                                image: base64,
                                audience: audience
                            })
                        });
                        
                        const data = await response.json();
                        resultDiv.textContent = JSON.stringify(data, null, 2);
                    } catch (error) {
                        resultDiv.textContent = 'Error: ' + error.message;
                    }
                };
                
                reader.readAsDataURL(fileInput.files[0]);
            }
        </script>
    </body>
    </html>
    """
    return html

# Mount Gradio at /gradio using WSGI middleware
from werkzeug.middleware.dispatcher import DispatcherMiddleware

app.wsgi_app = DispatcherMiddleware(
    app.wsgi_app,
    {'/gradio': gradio_app}
)

# RUN FLASK APP
if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    print("=" * 60)
    print(" Flask + Gradio Server Starting")
    print("=" * 60)
    print(f" Gradio UI:           http://localhost:{port}/gradio")
    print(f" Home (Redirect):     http://localhost:{port}/")
    print(f" API Test Page:       http://localhost:{port}/api/test")
    print(f" Health Check:        http://localhost:{port}/api/health")
    print(f" API Endpoint:        http://localhost:{port}/api/generate-story-base64")
    print("=" * 60)
    app.run(host="0.0.0.0", port=port, debug=False)