"""
routes/ai_summary.py — AI document analysis routes.
"""

import logging
import json
import io
import os
import re
import requests as http_requests
from flask import (Blueprint, render_template, request, redirect, url_for,
                   flash, session, jsonify)
from models import (
    get_all_criteria, get_active_criteria, create_criterion, update_criterion,
    delete_criterion, save_ai_analysis, get_ai_analysis_history,
    get_ai_analysis_detail, delete_ai_analysis, log_action,
)
from config import get_setting
from utils.decorators import login_required, admin_required

logger = logging.getLogger("routes.ai_summary")
ai_summary_bp = Blueprint("ai_summary", __name__, url_prefix="/ai-summary")

GROQ_MODELS = [
    "llama-3.3-70b-versatile",
    "llama-3.1-8b-instant",
    "gemma2-9b-it",
]

SUPPORTED_EXT = {".txt", ".md", ".csv", ".pdf", ".doc", ".docx", ".xlsx", ".xls"}


@ai_summary_bp.route("/")
@login_required
def ai_summary():
    user       = session["user"]
    criteria   = get_all_criteria()
    history    = get_ai_analysis_history(user_id=None)
    groq_key   = get_setting("groq.api_key", "")
    groq_model = get_setting("groq.model", GROQ_MODELS[0])
    return render_template("ai_summary.html",
                           criteria=criteria, history=history,
                           groq_key_set=bool(groq_key),
                           groq_model=groq_model,
                           groq_models=GROQ_MODELS,
                           is_admin=user["role"] == "admin")


@ai_summary_bp.route("/analyze", methods=["POST"])
@login_required
def analyze():
    user    = session["user"]
    files   = request.files.getlist("documents[]")
    model   = request.form.get("model", get_setting("groq.model", GROQ_MODELS[0]))
    # Read API key: env var takes priority, then DB setting
    api_key = (os.environ.get("GROQ_API_KEY") or "").strip() or get_setting("groq.api_key", "").strip()

    if not api_key:
        return jsonify({"error": "Groq API key is not configured. Set GROQ_API_KEY in .env or via Admin → Settings."}), 400
    if not files or all(f.filename == "" for f in files):
        return jsonify({"error": "No files uploaded."}), 400

    texts      = []
    file_names = []
    errors     = []

    for f in files:
        if not f.filename:
            continue
        ext = ("." + f.filename.rsplit(".", 1)[-1].lower()) if "." in f.filename else ""
        if ext not in SUPPORTED_EXT:
            errors.append(f"{f.filename}: unsupported file type")
            continue
        file_names.append(f.filename)
        try:
            content = _extract_text(f, ext)
            if content and content.strip():
                texts.append(f"=== {f.filename} ===\n{content.strip()}")
            else:
                errors.append(f"{f.filename}: no text could be extracted")
        except Exception as e:
            logger.warning(f"Could not extract text from {f.filename}: {e}")
            errors.append(f"{f.filename}: {e}")

    if not texts:
        detail = "; ".join(errors) if errors else "Check the file format and try again."
        return jsonify({"error": f"Could not extract text from any uploaded file. {detail}"}), 400

    criteria     = get_active_criteria()
    combined     = "\n\n".join(texts)

    try:
        result = _call_groq(api_key, model, combined, criteria)
    except Exception as e:
        logger.error(f"Groq API error: {e}")
        return jsonify({"error": f"AI analysis failed: {e}"}), 500

    criteria_snap = json.dumps([
        {"title": c["title"], "description": c["description"]} for c in criteria
    ])
    analysis_id = save_ai_analysis(
        user["id"], ", ".join(file_names), model,
        result.get("verdict"), criteria_snap, result.get("summary", ""),
        document_text=combined,
    )
    log_action(user["id"], "AI_ANALYSIS", "ai_analysis_log", analysis_id,
               f"AI analysis on {len(file_names)} file(s). Verdict: {result.get('verdict')}.")

    return jsonify({"analysis_id": analysis_id, "model": model,
                   "file_count": len(file_names), "file_errors": errors, **result})


# Magic-byte signatures for each supported binary format.
# Office Open XML (docx/xlsx) and legacy OLE (doc/xls) share these headers.
_FILE_MAGIC: dict = {
    ".pdf":  b"%PDF",
    ".docx": b"PK\x03\x04",
    ".xlsx": b"PK\x03\x04",
    ".doc":  b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
    ".xls":  b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
}


def _validate_magic(data: bytes, ext: str) -> None:
    """Raise ValueError if the file's actual bytes don't match its extension."""
    magic = _FILE_MAGIC.get(ext)
    if magic:
        if data[: len(magic)] != magic:
            raise ValueError(
                f"File content does not match the declared {ext} format "
                "(possible disguised upload)."
            )
    elif ext in (".txt", ".csv", ".md"):
        # Text files must be decodable; binary data masquerading as text is rejected.
        try:
            data[:512].decode("utf-8")
        except UnicodeDecodeError:
            raise ValueError(
                "Text file does not appear to be valid UTF-8. "
                "Ensure the file is a plain text document."
            )


def _extract_text(file_obj, ext: str) -> str:
    """Extract plain text from an uploaded file object."""
    data = file_obj.read()
    _validate_magic(data, ext)

    if ext in (".txt", ".md", ".csv"):
        return data.decode("utf-8", errors="replace")

    if ext == ".pdf":
        import pypdf
        reader = pypdf.PdfReader(io.BytesIO(data))
        pages  = []
        for page in reader.pages:
            t = page.extract_text()
            if t:
                pages.append(t)
        return "\n".join(pages)

    if ext in (".docx", ".doc"):
        # Use python-docx (installed as 'docx') — do NOT use docx2txt
        from docx import Document
        doc   = Document(io.BytesIO(data))
        lines = [para.text for para in doc.paragraphs if para.text.strip()]
        return "\n".join(lines)

    if ext in (".xlsx", ".xls"):
        import openpyxl
        wb    = openpyxl.load_workbook(io.BytesIO(data), read_only=True, data_only=True)
        lines = []
        for ws in wb.worksheets:
            for row in ws.iter_rows(values_only=True):
                line = "\t".join(str(c) if c is not None else "" for c in row)
                if line.strip():
                    lines.append(line)
        return "\n".join(lines)

    return ""


# Office address used as origin for distance/travel-time estimates.
_OFFICE_ADDRESS = "2815 Hartland Road, Falls Church, VA 22043, USA"

# Document text limit per API call.
# Stage 1 prompt template is ~1.8 KB overhead; 14 KB of doc text keeps the total
# JSON payload well under Groq's request size limit on all plan tiers.
_TEXT_LIMIT = 14_000

# Analyst persona — injected as the system message in every API call.
_SYSTEM_PROMPT = (
    "You are a senior government procurement analyst supporting a small business BD team. "
    "Extract and organize solicitation information with precision. "
    "Format all output in clean Markdown with tables and section headers. "
    "Write \"Not specified\" for any field not found in the documents — never guess."
)

# Stage 1 — extraction prompt (always sent). Kept compact to minimise payload size.
_EXTRACTION_PROMPT = """\
Analyze the {n} solicitation document(s) below as one combined package.
Fill in every table and section. Write **"Not specified"** for missing fields.

Our office: **{office}** — use as origin for travel estimates.

---

## Solicitation Overview

| Field | Value |
|---|---|
| Solicitation Number | |
| Solicitation Type | (RFP / RFQ / IFB / IDIQ / BPA / etc.) |
| Issuing Agency / Office | |
| Set-Aside | (Small Business / 8(a) / SDVOSB / HUBZone / WOSB / Unrestricted / etc.) |
| Contract Type | (FFP / T&M / Cost-Plus / IDIQ / etc.) |
| Contract Period | |
| NAICS Code | |
| Estimated Value | |

---

## Scope of Work

*(Detail required services, deliverables, and performance standards. Use bullet points.)*

---

## Performance Location(s)

*(All work sites; whether remote or on-site work is permitted.)*

---

## Key Dates & Deadlines

| Milestone | Date & Time |
|---|---|
| Pre-Proposal Conference / Site Visit | |
| Questions Due | |
| **Proposal Due** | |
| Award Date | |
| Period of Performance Start | |
| Other Deadlines | |

---

## Pre-Proposal Conference / Site Visit

| | |
|---|---|
| Date & Time | |
| Full Address | |
| Attendance | (Mandatory / Optional / N/A) |
| RSVP Required | |

---

## Point of Contact

| | |
|---|---|
| Name / Title | |
| Phone | |
| Email | |
| Questions Submission Method | |

---

## Proposal Requirements

*(Format, page limits, volumes, submission method, number of copies.)*

---

## Evaluation Factors

*(Government's evaluation criteria and weights as stated in the solicitation.)*

---

## Travel & Logistics

| | |
|---|---|
| Origin | {office} |
| Destination | |
| Estimated Driving Distance | |
| Estimated Drive Time | |

*(AI estimates — verify with a mapping service before scheduling.)*

---

## Red Flags & Notable Requirements

*(Security clearances, bonding, certifications, teaming restrictions, tight timelines, incumbent indicators, or other bid/no-bid risk factors.)*

---

## Total Square Footage

*(If applicable; write "N/A" if not.)*

---

## Overall Summary

- **What Is Being Procured:**
- **Contract Period & Value:**
- **Key Constraints:**
- **Immediate Action Items & Deadlines:**

---

DOCUMENTS:
{documents}"""

# Criteria evaluation — appended to the extraction prompt in the same API call.
# A single call avoids the rate-limit (429) that two back-to-back calls trigger
# on Groq's free tier.
_CRITERIA_SUFFIX = """

---

## Opportunity Alignment Evaluation

Now evaluate this opportunity against our company's criteria below.

**Our criteria:**
{criteria_list}

For each criterion provide:
- ✅ MEETS / ⚠️ PARTIALLY MEETS / ❌ DOES NOT MEET / ❓ CANNOT DETERMINE
- 1–2 sentences with specific evidence from the documents above.

**Overall Recommendation** — write exactly one of these lines (machine-read, do not alter):

  RECOMMENDATION: PURSUE
  RECOMMENDATION: PASS
  RECOMMENDATION: UNCLEAR

PURSUE = clearly aligns with most criteria and is competitive.
PASS   = fails one or more critical criteria or presents unacceptable risk.
UNCLEAR = insufficient information for a confident decision.

**Executive Summary:** 3–4 sentences in plain business language explaining the recommendation."""


def _call_groq(api_key: str, model: str, text: str, criteria: list) -> dict:
    """Single Groq API call: extraction + optional criteria evaluation in one request."""
    truncated = len(text) > _TEXT_LIMIT
    doc_text  = text[:_TEXT_LIMIT]
    n         = doc_text.count("=== ") or 1

    prompt = _EXTRACTION_PROMPT.format(
        n=n, office=_OFFICE_ADDRESS, documents=doc_text
    )

    if criteria:
        criteria_list = "\n".join(
            f"{i+1}. **{c['title']}**: {c['description']}"
            for i, c in enumerate(criteria)
        )
        prompt += _CRITERIA_SUFFIX.format(criteria_list=criteria_list)

    response = http_requests.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type":  "application/json",
        },
        json={
            "model":       model,
            "messages":    [
                {"role": "system", "content": _SYSTEM_PROMPT},
                {"role": "user",   "content": prompt},
            ],
            "max_tokens":  4096,
            "temperature": 0.1,
        },
        timeout=120,
    )
    if not response.ok:
        logger.error(f"Groq API {response.status_code}: {response.text[:300]}")
        response.raise_for_status()

    content = response.json()["choices"][0]["message"]["content"] or ""

    verdict = None
    if criteria:
        match = re.search(
            r"RECOMMENDATION\s*:\s*(PURSUE|PASS|UNCLEAR)",
            content, re.IGNORECASE,
        )
        if match:
            verdict = match.group(1).upper()

    return {"verdict": verdict, "summary": content, "truncated": truncated}


# ─── Criteria Management (admin only) ─────────────────────────────────────────

@ai_summary_bp.route("/criteria/create", methods=["POST"])
@admin_required
def create_criterion_view():
    admin      = session["user"]
    title      = request.form.get("title", "").strip()
    desc       = request.form.get("description", "").strip()
    is_active  = request.form.get("is_active", "1") == "1"
    try:
        sort_order = int(request.form.get("sort_order", 0) or 0)
    except (ValueError, TypeError):
        sort_order = 0
    try:
        create_criterion(admin["id"], title, desc, is_active, sort_order)
        flash(f"Criterion '{title}' created.", "success")
    except Exception as e:
        flash(f"Error: {e}", "danger")
    return redirect(url_for("ai_summary.ai_summary"))


@ai_summary_bp.route("/criteria/<int:criterion_id>/edit", methods=["POST"])
@admin_required
def edit_criterion_view(criterion_id):
    admin      = session["user"]
    title      = request.form.get("title", "").strip()
    desc       = request.form.get("description", "").strip()
    is_active  = request.form.get("is_active", "1") == "1"
    try:
        sort_order = int(request.form.get("sort_order", 0) or 0)
    except (ValueError, TypeError):
        sort_order = 0
    try:
        update_criterion(admin["id"], criterion_id, title, desc, is_active, sort_order)
        flash(f"Criterion '{title}' updated.", "success")
    except Exception as e:
        flash(f"Error: {e}", "danger")
    return redirect(url_for("ai_summary.ai_summary"))


@ai_summary_bp.route("/criteria/<int:criterion_id>/delete", methods=["POST"])
@admin_required
def delete_criterion_view(criterion_id):
    admin = session["user"]
    try:
        delete_criterion(admin["id"], criterion_id)
        flash("Criterion deleted.", "success")
    except Exception as e:
        flash(f"Error: {e}", "danger")
    return redirect(url_for("ai_summary.ai_summary"))


@ai_summary_bp.route("/chat", methods=["POST"])
@login_required
def chat():
    """Follow-up Q&A about a previously analyzed document."""
    user = session["user"]
    data = request.get_json(silent=True) or {}
    analysis_id = data.get("analysis_id")
    messages    = data.get("messages") or []

    if not analysis_id:
        return jsonify({"error": "analysis_id is required."}), 400
    if not messages or not isinstance(messages, list):
        return jsonify({"error": "messages is required."}), 400

    record = get_ai_analysis_detail(int(analysis_id))
    if not record:
        return jsonify({"error": "Analysis not found."}), 404

    api_key = (os.environ.get("GROQ_API_KEY") or "").strip() or get_setting("groq.api_key", "").strip()
    if not api_key:
        return jsonify({"error": "Groq API key is not configured."}), 400

    model = get_setting("groq.model", GROQ_MODELS[0])

    # Prefer raw document text so the AI can answer questions from source material.
    # Fall back to the stored summary for analyses run before this feature was added.
    raw_text = (record.get("document_text") or "").strip()
    if raw_text:
        context_label = "ORIGINAL DOCUMENT CONTENT"
        context_body  = raw_text[:12000]
    else:
        context_label = "ANALYZED DOCUMENT SUMMARY"
        context_body  = (record.get("summary_text") or "")[:8000]

    system_msg = (
        "You are a government procurement analyst assistant. "
        "The user has questions about a solicitation document. "
        "Answer based solely on the document content provided below. "
        "If the information is not present, say so clearly. Be concise and direct.\n\n"
        f"## {context_label}\n\n" + context_body
    )

    # Sanitize and cap conversation history at 20 turns
    api_messages = []
    for m in messages[-20:]:
        role    = m.get("role")    if isinstance(m, dict) else None
        content = (m.get("content") or "").strip() if isinstance(m, dict) else ""
        if role in ("user", "assistant") and content:
            api_messages.append({"role": role, "content": content})

    if not api_messages or api_messages[-1]["role"] != "user":
        return jsonify({"error": "Last message must be from the user."}), 400

    try:
        response = http_requests.post(
            "https://api.groq.com/openai/v1/chat/completions",
            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
            json={
                "model":       model,
                "messages":    [{"role": "system", "content": system_msg}] + api_messages,
                "max_tokens":  1024,
                "temperature": 0.2,
            },
            timeout=60,
        )
        if not response.ok:
            logger.error(f"Groq chat {response.status_code}: {response.text[:200]}")
            response.raise_for_status()
        reply = response.json()["choices"][0]["message"]["content"] or ""
        return jsonify({"reply": reply})
    except Exception as e:
        logger.error(f"AI chat error: {e}")
        return jsonify({"error": f"Chat failed: {e}"}), 500


@ai_summary_bp.route("/history/<int:analysis_id>")
@login_required
def analysis_detail(analysis_id):
    record = get_ai_analysis_detail(analysis_id)
    if not record:
        return jsonify({"error": "Not found"}), 404
    return jsonify({
        "id":                record["id"],
        "username":          record.get("username"),
        "file_names":        record["file_names"],
        "model":             record["model"],
        "verdict":           record["verdict"],
        "analyzed_at":       str(record["analyzed_at"]),
        "summary_text":      record["summary_text"],
        "criteria_snapshot": record.get("criteria_snapshot") or "",
    })


@ai_summary_bp.route("/history/<int:analysis_id>/delete", methods=["POST"])
@login_required
def delete_analysis(analysis_id):
    user = session["user"]
    if user["role"] != "admin":
        return jsonify({"error": "Access denied"}), 403
    record = get_ai_analysis_detail(analysis_id)
    if not record:
        return jsonify({"error": "Not found"}), 404
    delete_ai_analysis(analysis_id)
    log_action(user["id"], "DELETE_AI_ANALYSIS", "ai_analysis_log", analysis_id,
               f"Deleted AI analysis id={analysis_id}.")
    return jsonify({"ok": True})