app.py

import os
import logging
from flask import Flask, render_template, request, jsonify, session, redirect, url_for, flash
from flask_sqlalchemy import SQLAlchemy
from flask_login import LoginManager, login_user, logout_user, login_required, current_user
from sqlalchemy.orm import DeclarativeBase
import cohere
import requests
import uuid
from datetime import datetime

# Setup logging
logging.basicConfig(level=logging.INFO)

class Base(DeclarativeBase):
    pass

# ✅ Use writable instance_path to avoid PermissionError
app = Flask(__name__, instance_path='/tmp/instance')
app.secret_key = os.environ.get("SESSION_SECRET", os.urandom(24).hex())

# Database config
app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DATABASE_URL", "sqlite:///inkboard.db")
app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {
    'pool_pre_ping': True,
    "pool_recycle": 300,
}

db = SQLAlchemy(app, model_class=Base)

# Flask-Login setup
login_manager = LoginManager()
login_manager.init_app(app)
login_manager.login_view = 'login'
login_manager.login_message = 'Please log in to access InkBoard'

# API keys
API_KEY = os.environ.get("API_KEY")
HUGGINGFACE_API_KEY = os.environ.get("HUGGINGFACE_API_KEY")
HF_IMAGE_MODEL = "stabilityai/stable-diffusion-2-1"

cohere_client = cohere.Client(API_KEY) if API_KEY else None
logging.info(f"Received prompt: {scene_idea}")

@login_manager.user_loader
def load_user(user_id):
    from models import User
    return User.query.get(int(user_id))

@app.route('/')
def index():
    return render_template('dashboard.html', user=current_user) if current_user.is_authenticated else render_template('index.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    from models import User
    if request.method == 'POST':
        data = request.get_json() if request.is_json else request.form
        username, email, password = data.get('username', '').strip(), data.get('email', '').strip(), data.get('password', '').strip()

        if not all([username, email, password]):
            return jsonify({'error': 'All fields are required'}), 400

        if User.query.filter_by(username=username).first() or User.query.filter_by(email=email).first():
            return jsonify({'error': 'Username or email already exists'}), 400

        user = User(username=username, email=email)
        user.set_password(password)

        db.session.add(user)
        db.session.commit()
        login_user(user)

        return jsonify({'success': True, 'redirect': url_for('index')}) if request.is_json else redirect(url_for('index'))

    return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    from models import User
    if request.method == 'POST':
        data = request.get_json() if request.is_json else request.form
        username, password = data.get('username', '').strip(), data.get('password', '').strip()

        user = User.query.filter((User.username == username) | (User.email == username)).first()

        if user and user.check_password(password):
            login_user(user)
            return jsonify({'success': True, 'redirect': url_for('index')}) if request.is_json else redirect(url_for('index'))

        return jsonify({'error': 'Invalid credentials'}), 401

    return render_template('login.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('index'))


@app.route('/generate', methods=['POST'])
@login_required
def generate_content():
    from models import Creation

    # 1. ✅ Logging for debugging
    logging.info("⚡ generate_content() called")

    # 2. ✅ Extract prompt
    data = request.get_json()
    scene_idea = data.get('scene_idea', '').strip()
    logging.info(f"📥 Received prompt: {scene_idea}")

    if not scene_idea:
        logging.warning("❗ No scene idea provided")
        return jsonify({'error': 'Please provide a scene idea'}), 400

    if not cohere_client:
        logging.error("❌ Cohere client is not configured. Check your API_KEY.")
        return jsonify({'error': 'Text generation is unavailable'}), 500

    # 3. ✅ Try generating story
    try:
        story_prompt = f"""Transform this scene idea into a vivid paragraph:\nScene idea: {scene_idea}"""
        logging.info("🧠 Sending prompt to Cohere")

        story_response = cohere_client.generate(
            model='command',
            prompt=story_prompt,
            max_tokens=200,
            temperature=0.7,
            k=0
        )

        logging.info("✅ Cohere response received")

        expanded_story = story_response.generations[0].text.strip()
        logging.info(f"📝 Generated story: {expanded_story[:60]}...")

    except Exception as e:
        logging.error(f"❌ Cohere generation error: {str(e)}")
        return jsonify({'error': 'Failed to generate story'}), 500

    # 4. ✅ Generate image (optional)
    image_url = None
    try:
        image_url = generate_image_hf(scene_idea, expanded_story)
        logging.info(f"🖼️ Image generated: {image_url[:50]}...")
    except Exception as e:
        logging.warning(f"⚠️ Image generation failed: {e}")

    # 5. ✅ Save to DB
    creation_id = str(uuid.uuid4())
    creation = Creation(
        id=creation_id,
        user_id=current_user.id,
        scene_idea=scene_idea,
        story=expanded_story,
        image_url=image_url
    )
    db.session.add(creation)
    db.session.commit()

    return jsonify({
        'success': True,
        'story': expanded_story,
        'image_url': image_url,
        'creation_id': creation_id
    })


@app.route('/save_journal', methods=['POST'])
@login_required
def save_journal():
    from models import Creation
    data = request.get_json()
    creation_id, journal_entry = data.get('creation_id'), data.get('journal_entry', '').strip()
    creation = Creation.query.filter_by(id=creation_id, user_id=current_user.id).first()
    if not creation:
        return jsonify({'error': 'Not found'}), 404
    creation.journal_entry = journal_entry
    creation.updated_at = datetime.utcnow()
    db.session.commit()
    return jsonify({'success': True})

@app.route('/get_creations')
@login_required
def get_creations():
    from models import Creation
    creations = Creation.query.filter_by(user_id=current_user.id).order_by(Creation.created_at.desc()).all()
    return jsonify({'creations': [{
        'id': c.id,
        'scene_idea': c.scene_idea,
        'story': c.story,
        'image_url': c.image_url,
        'journal_entry': c.journal_entry,
        'created_at': c.created_at.isoformat()
    } for c in creations]})

def generate_image_hf(scene_idea, story):
    prompt = f"An artistic illustration of: {scene_idea}, dreamy and vivid"
    models = [
        "runwayml/stable-diffusion-v1-5",
        "stabilityai/stable-diffusion-2-1",
        "CompVis/stable-diffusion-v1-4"
    ]

    for model in models:
        try:
            logging.info(f"Trying model: {model}")
            res = requests.post(
                f"https://api-inference.huggingface.co/models/{model}",
                headers={"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"},
                json={"inputs": prompt},
                timeout=15  # 15 seconds max
            )

            logging.info(f"Status: {res.status_code}")
            if res.status_code == 200:
                import base64
                return f"data:image/png;base64,{base64.b64encode(res.content).decode()}"

            else:
                logging.warning(f"Model {model} failed: {res.status_code} - {res.text[:200]}")

        except Exception as e:
            logging.error(f"Error from model {model}: {e}")
            continue

    return generate_svg_placeholder(scene_idea, story)
        
# DB setup
from models import User, Creation
with app.app_context():
    db.create_all()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)