from flask import Blueprint, render_template, request, redirect, url_for, flash, send_file
from flask_login import login_required, current_user
from app.extensions import db
from app.models.features import LegalDocument
from app.models.report import Report

legal_bp = Blueprint("legal", __name__, url_prefix="/legal")

DOC_TYPES = [
    ("incorporation", "عقد تأسيس شركة", "Company Incorporation Agreement"),
    ("partnership", "اتفاقية شراكة", "Partnership Agreement"),
    ("nda", "اتفاقية عدم إفشاء (NDA)", "Non-Disclosure Agreement"),
    ("employment", "عقد عمل", "Employment Contract"),
    ("shareholders", "اتفاقية المساهمين", "Shareholders Agreement"),
    ("terms", "شروط الخدمة", "Terms of Service"),
    ("privacy", "سياسة الخصوصية", "Privacy Policy"),
    ("investment", "اتفاقية استثمار", "Investment Agreement"),
]


@legal_bp.route("/")
@login_required
def index():
    docs = LegalDocument.query.filter_by(user_id=current_user.id).order_by(LegalDocument.created_at.desc()).all()
    reports = Report.query.filter_by(user_id=current_user.id).order_by(Report.created_at.desc()).all()
    return render_template("legal/index.html", docs=docs, doc_types=DOC_TYPES, reports=reports)


@legal_bp.route("/generate", methods=["POST"])
@login_required
def generate():
    doc_type = request.form.get("doc_type", "").strip()
    report_id = request.form.get("report_id", type=int)
    lang = request.form.get("language", "ar")

    type_info = next((t for t in DOC_TYPES if t[0] == doc_type), None)
    if not type_info:
        flash("Invalid document type.", "error")
        return redirect(url_for("legal.index"))

    report = Report.query.get(report_id) if report_id else None
    project_context = ""
    if report and report.user_id == current_user.id:
        project_context = f"\nProject: {report.project_name}\nDescription: {report.project_description}\nSynthesis: {str(report.synthesis_result)[:2000]}"

    try:
        from app.ai.diamond import call_ai, run_sync
        prompt = f"""Generate a professional {type_info[2]} document for a Saudi Arabian company.
{project_context}
Requirements:
- Follow Saudi Commercial Law and Companies Law
- Include all standard clauses
- Be thorough and legally sound
- Language: {"Arabic" if lang == "ar" else "English"}
- Format with clear sections and numbered clauses
- Include blanks for: company name, date, party names, amounts where applicable"""

        messages = [{"role": "system", "content": "You are a Saudi legal expert. Generate complete, professional legal documents."}, {"role": "user", "content": prompt}]
        content = run_sync(call_ai(messages, model_type="deep"))

        doc = LegalDocument(user_id=current_user.id, report_id=report_id, doc_type=doc_type,
                            title=type_info[1], content=content, language=lang)
        db.session.add(doc)
        db.session.commit()
        return redirect(url_for("legal.view", doc_id=doc.id))
    except Exception as e:
        flash(f"Error generating document: {str(e)}", "error")
        return redirect(url_for("legal.index"))


@legal_bp.route("/view/<int:doc_id>")
@login_required
def view(doc_id):
    doc = LegalDocument.query.get_or_404(doc_id)
    if doc.user_id != current_user.id:
        flash("Access denied.", "error")
        return redirect(url_for("legal.index"))
    return render_template("legal/view.html", doc=doc)


@legal_bp.route("/export/<int:doc_id>")
@login_required
def export_docx(doc_id):
    doc = LegalDocument.query.get_or_404(doc_id)
    if doc.user_id != current_user.id:
        flash("Access denied.", "error")
        return redirect(url_for("legal.index"))
    try:
        from docx import Document as DocxDoc
        from docx.shared import Pt
        import io
        docx = DocxDoc()
        docx.add_heading(doc.title, 0)
        for para in doc.content.split("\n"):
            if para.strip():
                docx.add_paragraph(para.strip())
        buf = io.BytesIO()
        docx.save(buf)
        buf.seek(0)
        return send_file(buf, as_attachment=True, download_name=f"jadwa_legal_{doc.doc_type}.docx",
                         mimetype="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
    except ImportError:
        flash("python-docx required.", "error")
        return redirect(url_for("legal.view", doc_id=doc_id))