# src/routes/analysis.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.analysis import ImpactAnalysis
from src.models.legislation import Legislation
from src.models.draft import Draft
from src.models.amendment import Amendment
from src.utils.groq_client import perform_impact_analysis # Import Groq utility

analysis_bp = Blueprint("analysis", __name__, url_prefix="/analysis", template_folder="../templates/analysis")

@analysis_bp.route("/")
@login_required
def list_analyses():
    """Lists all previously generated impact analyses."""
    analyses = ImpactAnalysis.query.order_by(ImpactAnalysis.generated_at.desc()).all()
    return render_template("list_analyses.html", analyses=analyses)

@analysis_bp.route("/run", methods=["GET", "POST"])
@login_required
def run_analysis():
    """Form to select document and run a new impact analysis using Groq."""
    if request.method == "POST":
        document_type = request.form.get("document_type")
        document_id = request.form.get("document_id", type=int)
        analysis_type = request.form.get("analysis_type")

        if not document_type or not document_id or not analysis_type:
            flash("Document Type, Document ID, and Analysis Type are required.", "warning")
            legislations = Legislation.query.order_by(Legislation.title).all()
            drafts = Draft.query.order_by(Draft.title).all()
            amendments = Amendment.query.order_by(Amendment.id).all()
            return render_template("run_analysis.html", legislations=legislations, drafts=drafts, amendments=amendments, 
                                   document_type=document_type, document_id=document_id, analysis_type=analysis_type)

        # Fetch the document content
        document_content = None
        document_title = ""
        doc = None
        if document_type == "Legislation":
            doc = Legislation.query.get(document_id)
            if doc: 
                document_content = doc.full_text
                document_title = doc.title
        elif document_type == "Draft":
            doc = Draft.query.get(document_id)
            if doc: 
                document_content = doc.content
                document_title = doc.title
        elif document_type == "Amendment":
            doc = Amendment.query.get(document_id)
            if doc: 
                document_content = doc.proposed_changes
                document_title = f'Amendment #{doc.id} for {doc.legislation.title if doc.legislation else "Unknown"}'

        if not doc:
             flash(f"Could not find {document_type} with ID {document_id}.", "danger")
             legislations = Legislation.query.order_by(Legislation.title).all()
             drafts = Draft.query.order_by(Draft.title).all()
             amendments = Amendment.query.order_by(Amendment.id).all()
             return render_template("run_analysis.html", legislations=legislations, drafts=drafts, amendments=amendments)

        if not document_content:
            flash(f"Document content for {document_type} ID {document_id} is empty. Cannot run analysis.", "warning")
            return redirect(url_for("analysis.run_analysis"))

        # Call Groq for impact analysis
        flash(f"Running {analysis_type} impact analysis using AI... This may take a moment.", "info")
        predicted_impact, confidence_score, rationale = perform_impact_analysis(document_content, analysis_type)

        if predicted_impact is None or predicted_impact.startswith("Error:"):
            flash(f'AI impact analysis failed: {predicted_impact or "No response from model."}', "danger")
            # Optionally save a failed analysis record or just redirect
            return redirect(url_for("analysis.run_analysis"))
        
        # Save the successful analysis result
        new_analysis = ImpactAnalysis(
            document_id=document_id,
            document_type=document_type,
            analysis_type=analysis_type,
            predicted_impact=predicted_impact,
            confidence_score=confidence_score, # Can be None if parsing failed
            rationale=rationale, # Can be None if parsing failed
            generated_by_user_id=session.get("user_id")
        )

        try:
            db.session.add(new_analysis)
            db.session.commit()
            flash(f"{analysis_type} impact analysis for 	'{document_title}' completed successfully.", "success")
            return redirect(url_for("analysis.view_analysis", analysis_id=new_analysis.id))
        except Exception as e:
            db.session.rollback()
            flash(f"Error saving analysis result: {e}", "danger")
            return redirect(url_for("analysis.list_analyses"))

    # GET request: Show the form
    legislations = Legislation.query.order_by(Legislation.title).all()
    drafts = Draft.query.order_by(Draft.title).all()
    amendments = Amendment.query.order_by(Amendment.id).all()
    return render_template("run_analysis.html", legislations=legislations, drafts=drafts, amendments=amendments)

@analysis_bp.route("/<int:analysis_id>")
@login_required
def view_analysis(analysis_id):
    """Displays the results of a specific impact analysis."""
    analysis = ImpactAnalysis.query.get_or_404(analysis_id)
    document_title = "Unknown Document"
    if analysis.document_type == "Legislation":
        doc = Legislation.query.get(analysis.document_id)
        if doc: document_title = doc.title
    elif analysis.document_type == "Draft":
        doc = Draft.query.get(analysis.document_id)
        if doc: document_title = doc.title
    elif analysis.document_type == "Amendment":
        doc = Amendment.query.get(analysis.document_id)
        if doc: document_title = f'Amendment #{doc.id} for {doc.legislation.title if doc.legislation else "Unknown"}'
        
    return render_template("view_analysis.html", analysis=analysis, document_title=document_title)