src/routes/drafting.py

# src/routes/drafting.py

from flask import Blueprint, render_template, request, redirect, url_for, flash, session
from src.extensions import db
from src.decorators import login_required
from src.models.draft import Draft
from src.models.legislation import Legislation # Needed for linking drafts
from src.utils.groq_client import generate_draft_content # Import Groq utility

drafting_bp = Blueprint("drafting", __name__, url_prefix="/drafts", template_folder="../templates/drafts")

@drafting_bp.route("/")
@login_required
def list_drafts():
    """Lists all current drafts."""
    drafts = Draft.query.order_by(Draft.last_updated.desc()).all()
    return render_template("list_drafts.html", drafts=drafts)

@drafting_bp.route("/new", methods=["GET", "POST"])
@login_required
def create_draft():
    """Handles creation of a new draft, optionally using Groq for initial content."""
    if request.method == "POST":
        title = request.form.get("title")
        policy_intent = request.form.get("policy_intent")
        initial_content = request.form.get("initial_content", "")
        related_legislation_id = request.form.get("related_legislation_id")
        use_ai_generation = "use_ai_generation" in request.form

        if not title or not policy_intent:
            flash("Title and Policy Intent are required.", "warning")
            legislations = Legislation.query.order_by(Legislation.title).all()
            return render_template("create_draft.html", legislations=legislations, title=title, policy_intent=policy_intent, initial_content=initial_content, related_legislation_id=related_legislation_id, use_ai_generation=use_ai_generation)

        author_id = session.get("user_id")
        generated_content = None

        # Generate content using Groq if requested and no initial content provided
        if use_ai_generation and not initial_content:
            flash("Attempting to generate initial draft content using AI...", "info")
            generated_content = generate_draft_content(policy_intent)
            if generated_content and not generated_content.startswith("Error:"):
                initial_content = generated_content
                flash("AI successfully generated initial draft content.", "success")
            elif generated_content:
                flash(f"AI generation failed: {generated_content}", "danger")
            else:
                 flash("AI generation failed: No response from model.", "danger")

        new_draft = Draft(
            title=title,
            policy_intent=policy_intent,
            content=initial_content,
            author_id=author_id,
            related_legislation_id=int(related_legislation_id) if related_legislation_id else None
        )

        try:
            db.session.add(new_draft)
            db.session.commit()
            flash(f"Draft 	'{title}' created successfully.", "success")
            return redirect(url_for("drafting.view_draft", draft_id=new_draft.id))
        except Exception as e:
            db.session.rollback()
            flash(f"Error creating draft: {e}", "danger")
            # Re-render form on error
            legislations = Legislation.query.order_by(Legislation.title).all()
            return render_template("create_draft.html", legislations=legislations, title=title, policy_intent=policy_intent, initial_content=initial_content, related_legislation_id=related_legislation_id, use_ai_generation=use_ai_generation)

    # GET request: Show the form
    legislations = Legislation.query.order_by(Legislation.title).all()
    return render_template("create_draft.html", legislations=legislations)

@drafting_bp.route("/<int:draft_id>")
@login_required
def view_draft(draft_id):
    """Displays a single draft for viewing or editing."""
    draft = Draft.query.get_or_404(draft_id)
    return render_template("view_draft.html", draft=draft)

# Add routes for updating and deleting drafts later if required