from __future__ import annotations

import shutil
import threading
from datetime import timedelta
from pathlib import Path

import psutil
import pymupdf
from flask import (
    Flask,
    jsonify,
    redirect,
    render_template,
    request,
    send_file,
    url_for,
)
from werkzeug.middleware.proxy_fix import ProxyFix
from werkzeug.utils import secure_filename

import config
from rag.auth import AtlasAuth
from rag.extractors import DocumentExtractor, ExtractionOptions
from rag.groq_service import GroqService
from rag.ingestion import IngestionManager, quick_fingerprint
from rag.models import EmbeddingService, RerankerService
from rag.retrieval import RetrievalService
from rag.state import StateDB
from rag.store import VectorStore
from rag.uploads import UploadError, UploadManager

app = Flask(__name__)
app.secret_key = config.FLASK_SECRET_KEY
app.config.update(
    MAX_CONTENT_LENGTH=config.UPLOAD_CHUNK_BYTES + 1024 * 1024,
    # Nouveau nom pour éviter qu'une ancienne cookie SameSite=Lax interfère.
    SESSION_COOKIE_NAME="atlas_session_v2",
    SESSION_COOKIE_HTTPONLY=True,
    SESSION_COOKIE_SECURE=config.IS_HF_SPACE,
    # Dans l'onglet intégré Hugging Face, Atlas tourne dans une iframe
    # cross-site (huggingface.co -> *.hf.space). La cookie OAuth doit donc
    # être utilisable dans ce contexte.
    SESSION_COOKIE_SAMESITE="None" if config.IS_HF_SPACE else "Lax",
    SESSION_COOKIE_PARTITIONED=config.IS_HF_SPACE,
    PERMANENT_SESSION_LIFETIME=timedelta(hours=8),
)
app.wsgi_app = ProxyFix(
    app.wsgi_app,
    x_for=1,
    x_proto=1,
    x_host=1,
    x_prefix=1,
)

auth = AtlasAuth(app, config)
manager_required = auth.require_roles("manager", "admin")
admin_required = auth.require_roles("admin")

state = StateDB(config.STATE_DB_PATH)
embedder = EmbeddingService(
    config.EMBEDDING_MODEL, config.EMBEDDING_DIMENSION,
    config.EMBED_BATCH_SIZE, config.TORCH_NUM_THREADS,
)
reranker = RerankerService(
    config.RERANKER_MODEL, config.RERANK_BATCH_SIZE,
    config.RERANK_ENABLED, config.TORCH_NUM_THREADS,
)
store = VectorStore(
    config.LANCEDB_DIR,
    config.EMBEDDING_DIMENSION,
    config.INDEX_STATE_PATH,
    records_dir=config.CHUNK_RECORDS_DIR,
    legacy_db_dir=config.LEGACY_LANCEDB_DIR,
)


def restore_catalog_from_persistent_data() -> None:
    """Rend immédiatement utilisables les PDF et LanceDB du Bucket.

    Le Bucket initial contient déjà les PDF et les vecteurs, mais pas forcément
    le fichier SQLite local. On recrée uniquement le catalogue; aucun embedding
    et aucune réindexation ne sont recalculés.
    """
    summaries = store.document_summaries()
    if not summaries:
        return

    files_by_name = {
        path.name: path
        for path in config.DOCUMENTS_DIR.rglob("*")
        if path.is_file()
        and path.suffix.lower() in config.SUPPORTED_EXTENSIONS
    }

    restored = 0
    removed_orphans: list[str] = []
    cleanup_errors: list[str] = []
    for summary in summaries:
        document_id = str(summary.get("document_id") or "")
        filename = str(summary.get("filename") or "")
        source = files_by_name.get(filename)
        if source is None:
            label = filename or document_id
            try:
                # Le PDF a été supprimé mais d'anciens fragments peuvent être
                # restés dans la LanceDB historique après l'erreur FUSE. Ils ne
                # doivent jamais revenir dans les réponses du RAG.
                store.delete_document(document_id)
                state.delete_document(document_id)
                removed_orphans.append(label)
            except Exception as exc:
                cleanup_errors.append(
                    f"{label} ({type(exc).__name__}: {exc})"
                )
            continue

        state.restore_completed_document(
            document_id=str(summary["document_id"]),
            fingerprint=quick_fingerprint(source),
            filename=filename,
            path=str(source.resolve()),
            extension=source.suffix.lower(),
            total_pages=int(summary.get("total_pages") or 0),
            chunks=int(summary.get("chunks") or 0),
        )
        restored += 1

    if restored:
        print(
            f"[Atlas] Catalogue restauré : {restored} document(s) depuis le Bucket.",
            flush=True,
        )
    if removed_orphans:
        print(
            "[Atlas] Fragments orphelins supprimés : "
            + ", ".join(removed_orphans),
            flush=True,
        )
    if cleanup_errors:
        print(
            "[Atlas] Échec du nettoyage de fragments orphelins : "
            + "; ".join(cleanup_errors),
            flush=True,
        )


restore_catalog_from_persistent_data()

extractor = DocumentExtractor(ExtractionOptions(
    ocr_enabled=config.OCR_ENABLED,
    ocr_language=config.OCR_LANGUAGE,
    ocr_dpi=config.OCR_DPI,
    ocr_min_text_chars=config.OCR_MIN_TEXT_CHARS,
    tesseract_cmd=config.TESSERACT_CMD,
    extract_tables=config.EXTRACT_TABLES,
))
ingestion = IngestionManager(
    state, store, embedder, extractor,
    config.PAGE_BATCH_SIZE, config.CHUNK_WORDS, config.CHUNK_OVERLAP_WORDS,
    config.SUPPORTED_EXTENSIONS, config.AUTO_OPTIMIZE,
)
retrieval = RetrievalService(
    store, embedder, reranker,
    config.DENSE_CANDIDATES, config.TEXT_CANDIDATES, config.RRF_K,
    config.RERANK_POOL, config.FINAL_SOURCES, config.NEIGHBOR_PAGES,
)
groq_service = GroqService(config.GROQ_API_KEY, config.GROQ_MODEL, config.MAX_CONTEXT_CHARS)
index_lock = threading.Lock()
uploads = UploadManager(
    config.STATE_DB_PATH,
    config.UPLOAD_TEMP_DIR,
    config.UPLOADS_DIR,
    config.SUPPORTED_EXTENSIONS,
    config.UPLOAD_CHUNK_BYTES,
    config.UPLOAD_RESERVE_BYTES,
)


def machine_stats() -> dict:
    memory = psutil.virtual_memory()
    disk = shutil.disk_usage(config.DATA_DIR)
    return {
        "ram_total_gb": round(memory.total / 1024**3, 1),
        "ram_used_gb": round(memory.used / 1024**3, 1),
        "ram_percent": round(memory.percent, 1),
        "disk_free_gb": round(disk.free / 1024**3, 1),
        "disk_total_gb": round(disk.total / 1024**3, 1),
    }


PUBLIC_ENDPOINTS = {
    "login",
    "oauth_callback",
    "logout",
    "api_health",
    "static",
}
UNSAFE_METHODS = {"POST", "PUT", "PATCH", "DELETE"}


@app.before_request
def enforce_authentication_and_csrf():
    if request.endpoint in PUBLIC_ENDPOINTS:
        return None

    user = auth.current_user()
    if not user:
        if request.path.startswith("/api/"):
            return jsonify({"error": "Authentification requise."}), 401
        return redirect(url_for("login", next=request.full_path))

    if request.method in UNSAFE_METHODS and not auth.csrf_valid():
        return jsonify({
            "error": "Jeton de sécurité invalide. Recharge la page puis réessaie."
        }), 403

    return None


@app.get("/login")
def login():
    return auth.login_redirect(request.args.get("next"))


@app.get("/auth/callback")
def oauth_callback():
    try:
        destination = auth.finish_login()
        return redirect(destination)
    except PermissionError as exc:
        return (
            "<h1>Accès refusé</h1>"
            f"<p>{str(exc)}</p>"
            "<p>Demande à l'administrateur Atlas de vérifier ton compte.</p>",
            403,
        )
    except Exception as exc:
        app.logger.exception("Échec OAuth Hugging Face")
        return (
            "<h1>Connexion impossible</h1>"
            "<p>La connexion Hugging Face a échoué. "
            "Recharge le Space et réessaie.</p>"
            f"<p>Détail : {type(exc).__name__}</p>",
            500,
        )


@app.get("/logout")
def logout():
    auth.logout()
    return render_template("logged_out.html")


@app.get("/")
def home():
    return render_template("index.html")


@app.get("/api/me")
def api_me():
    user = auth.current_user()
    if not user:
        return jsonify({"error": "Authentification requise."}), 401
    return jsonify({
        "authenticated": True,
        "user": user,
        "csrfToken": auth.csrf_token(),
    })


@app.get("/api/config")
def api_config():
    return jsonify({
        "groqConfigured": groq_service.configured,
        "groqModel": config.GROQ_MODEL,
        "embeddingModel": config.EMBEDDING_MODEL,
        "rerankerModel": config.RERANKER_MODEL,
        "rerankerEnabled": config.RERANK_ENABLED,
        "queryExpansion": config.QUERY_EXPANSION_ENABLED,
        "ocrEnabled": config.OCR_ENABLED,
        "supportedExtensions": sorted(config.SUPPORTED_EXTENSIONS),
        "uploadChunkBytes": config.UPLOAD_CHUNK_BYTES,
        "uploadChunkMB": config.UPLOAD_CHUNK_MB,
    })


@app.get("/api/models")
def api_models():
    try:
        return jsonify({"models": groq_service.list_models()})
    except Exception as exc:
        return jsonify({"models": [config.GROQ_MODEL], "warning": str(exc)})


@app.get("/api/status")
@manager_required
def api_status():
    return jsonify({
        "stats": {**state.stats(), "vector_rows": store.count()},
        "machine": machine_stats(),
        "ingestion": {
            "paused": ingestion.paused,
            "queueSize": ingestion.queue_size,
            "currentDocumentId": ingestion.current_document_id,
        },
        "models": {"embeddingLoaded": embedder.loaded, "rerankerLoaded": reranker.loaded},
        "documents": state.list_documents(),
        "uploads": uploads.list_recent(50),
    })


@app.get("/api/events")
@manager_required
def api_events():
    return jsonify({"events": state.list_events(int(request.args.get("limit", 100)))})


@app.post("/api/uploads/start")
@manager_required
def api_upload_start():
    payload = request.get_json(silent=True) or {}
    filename = secure_filename(str(payload.get("filename", "")))
    try:
        session = uploads.start(
            filename=filename,
            total_size=int(payload.get("size", 0)),
            fingerprint=str(payload.get("fingerprint", "")),
        )
        return jsonify({"upload": session, "chunkSize": config.UPLOAD_CHUNK_BYTES})
    except (UploadError, ValueError, TypeError) as exc:
        return jsonify({"error": str(exc)}), 400


@app.get("/api/uploads/<upload_id>")
@manager_required
def api_upload_status(upload_id: str):
    session = uploads.get(upload_id)
    if not session:
        return jsonify({"error": "Session d’upload inconnue."}), 404
    return jsonify({"upload": session, "chunkSize": config.UPLOAD_CHUNK_BYTES})


@app.post("/api/uploads/<upload_id>/chunk")
@manager_required
def api_upload_chunk(upload_id: str):
    try:
        offset = int(request.headers.get("X-Upload-Offset", "-1"))
        block = request.get_data(cache=False, as_text=False)
        session = uploads.append(upload_id, offset, block)
        return jsonify({"upload": session})
    except (UploadError, ValueError, TypeError) as exc:
        session = uploads.get(upload_id)
        status = 409 if session else 404
        return jsonify({"error": str(exc), "upload": session}), status


@app.post("/api/uploads/<upload_id>/complete")
@manager_required
def api_upload_complete(upload_id: str):
    try:
        session, path = uploads.complete(upload_id)
        documents = ingestion.enqueue_paths([path])
        return jsonify({
            "upload": session,
            "documents": documents,
            "message": "Upload terminé. Le document a été ajouté à la file d’indexation.",
        })
    except UploadError as exc:
        return jsonify({"error": str(exc), "upload": uploads.get(upload_id)}), 409


@app.delete("/api/uploads/<upload_id>")
@manager_required
def api_upload_cancel(upload_id: str):
    try:
        session = uploads.cancel(upload_id)
        return jsonify({"upload": session, "message": "Upload annulé et fichier temporaire supprimé."})
    except UploadError as exc:
        return jsonify({"error": str(exc)}), 409


@app.post("/api/ingest-folder")
@admin_required
def api_ingest_folder():
    paths = [
        path for path in config.DOCUMENTS_DIR.rglob("*")
        if path.is_file() and path.suffix.lower() in config.SUPPORTED_EXTENSIONS
    ]
    return jsonify({"found": len(paths), "documents": ingestion.enqueue_paths(paths)})


@app.post("/api/ingestion/pause")
@admin_required
def api_pause():
    ingestion.pause()
    return jsonify({"message": "Indexation en pause après le lot en cours."})


@app.post("/api/ingestion/resume")
@admin_required
def api_resume():
    ingestion.resume()
    return jsonify({"message": "Indexation reprise."})


@app.post("/api/documents/<document_id>/cancel")
@manager_required
def api_cancel(document_id: str):
    ingestion.cancel(document_id)
    return jsonify({"message": "Annulation demandée."})


@app.post("/api/documents/<document_id>/retry")
@manager_required
def api_retry(document_id: str):
    if not ingestion.retry(document_id):
        return jsonify({"error": "Document introuvable ou déjà dans la file."}), 404
    return jsonify({"message": "Document remis dans la file."})


@app.post("/api/documents/<document_id>/reindex")
@manager_required
def api_reindex(document_id: str):
    doc = state.get_document(document_id)
    if not doc:
        return jsonify({"error": "Document inconnu."}), 404
    if ingestion.is_active(document_id):
        return jsonify({"error": "Ce document est encore actif. Annule-le et attends l’état Annulé avant de le réindexer."}), 409
    documents = ingestion.enqueue_paths([Path(doc["path"])], force=True)
    return jsonify({"message": "Réindexation lancée.", "documents": documents})


def _managed_document_path(raw_path: str) -> Path:
    """Retourne un chemin sûr appartenant au stockage géré par Atlas."""
    path = Path(raw_path).expanduser().resolve()
    allowed_roots = {
        Path(config.DOCUMENTS_DIR).expanduser().resolve(),
        Path(config.UPLOADS_DIR).expanduser().resolve(),
    }

    if not any(path == root or path.is_relative_to(root) for root in allowed_roots):
        raise ValueError(
            "Le fichier est en dehors des dossiers gérés par Atlas; "
            "suppression physique refusée."
        )
    return path


@app.delete("/api/documents/<document_id>")
@manager_required
def api_delete(document_id: str):
    if ingestion.is_active(document_id):
        return jsonify({
            "error": (
                "Ce document est encore dans la file ou en cours. "
                "Annule-le d’abord, puis attends l’état Annulé."
            )
        }), 409

    doc = state.get_document(document_id)
    if not doc:
        return jsonify({"error": "Document inconnu."}), 404

    try:
        source_path = _managed_document_path(str(doc["path"]))
    except (KeyError, TypeError, ValueError) as exc:
        return jsonify({"error": str(exc)}), 400

    # On nettoie d'abord la base active locale et l'archive persistante.
    # Ainsi, si la suppression du PDF échoue, l'utilisateur peut simplement
    # relancer l'action sans retrouver des fragments fantômes.
    try:
        store.delete_document(document_id)
    except Exception as exc:
        return jsonify({
            "error": (
                "Impossible de nettoyer les fragments du document : "
                f"{type(exc).__name__}: {exc}"
            )
        }), 500

    try:
        if source_path.exists():
            if not source_path.is_file():
                return jsonify({
                    "error": "Le chemin du document n’est pas un fichier."
                }), 409
            source_path.unlink()
    except OSError as exc:
        return jsonify({
            "error": (
                "Les fragments ont été supprimés, mais le PDF n’a pas pu "
                f"être effacé. Relance Supprimer : {exc}"
            )
        }), 500

    try:
        state.delete_document(document_id)
    except Exception as exc:
        return jsonify({
            "error": (
                "Le fichier et les fragments ont été supprimés, mais "
                "l’entrée du catalogue n’a pas pu être nettoyée. "
                f"Relance Supprimer : {type(exc).__name__}: {exc}"
            )
        }), 500

    return jsonify({
        "message": (
            "Document supprimé définitivement : fichier, fragments persistants "
            "et état d’indexation effacés."
        )
    })


@app.get("/api/documents/<document_id>/file")
def api_document_file(document_id: str):
    doc = state.get_document(document_id)
    if not doc or not Path(doc["path"]).is_file():
        return jsonify({"error": "Fichier introuvable."}), 404
    return send_file(doc["path"], as_attachment=False, download_name=doc["filename"])


@app.get("/api/documents/<document_id>/page/<int:page_number>/preview")
def api_page_preview(document_id: str, page_number: int):
    doc = state.get_document(document_id)
    if not doc or Path(doc["path"]).suffix.lower() != ".pdf":
        return jsonify({"error": "Aperçu disponible uniquement pour les PDF."}), 404
    path = Path(doc["path"])
    if not path.is_file():
        return jsonify({"error": "Fichier introuvable."}), 404
    with pymupdf.open(path) as pdf:
        if page_number < 1 or page_number > pdf.page_count:
            return jsonify({"error": "Page invalide."}), 404
        page = pdf.load_page(page_number - 1)
        pix = page.get_pixmap(matrix=pymupdf.Matrix(1.4, 1.4), alpha=False)
        return app.response_class(pix.tobytes("png"), mimetype="image/png")


def run_retrieval(question: str, payload: dict):
    expansion = bool(payload.get("queryExpansion", config.QUERY_EXPANSION_ENABLED))
    model = str(payload.get("model", "")).strip() or None
    variants = groq_service.expand_queries(question, enabled=expansion, model=model)
    return retrieval.search(question, variants)


@app.post("/api/search")
def api_search():
    payload = request.get_json(silent=True) or {}
    question = str(payload.get("query", "")).strip()
    if not question:
        return jsonify({"error": "Écris une question."}), 400
    if store.count() == 0:
        return jsonify({"error": "La base est vide. Indexe d’abord un document."}), 400
    try:
        result = run_retrieval(question, payload)
        return jsonify({"query": question, **result})
    except Exception as exc:
        return jsonify({"error": f"Recherche impossible : {type(exc).__name__}: {exc}"}), 500


@app.post("/api/ask")
def api_ask():
    payload = request.get_json(silent=True) or {}
    question = str(payload.get("question", "")).strip()
    model = str(payload.get("model", "")).strip() or None
    if not question:
        return jsonify({"error": "Écris une question."}), 400
    if store.count() == 0:
        return jsonify({"error": "La base est vide. Indexe d’abord un document."}), 400
    try:
        result = run_retrieval(question, payload)
        answer = groq_service.answer(question, result["sources"], model=model)
        return jsonify({"question": question, "answer": answer, **result})
    except Exception as exc:
        return jsonify({"error": f"Réponse impossible : {type(exc).__name__}: {exc}"}), 500


@app.post("/api/indexes/optimize")
@admin_required
def api_optimize():
    if not index_lock.acquire(blocking=False):
        return jsonify({"error": "Une optimisation est déjà en cours."}), 409

    def worker():
        try:
            result = store.build_indexes(full_rebuild=True)
            state.add_event(result.get("message", "Index optimisés."))
        except Exception as exc:
            state.add_event(f"Échec optimisation : {exc}", level="error")
        finally:
            index_lock.release()

    threading.Thread(target=worker, daemon=True, name="atlas-index-builder").start()
    return jsonify({"message": "Optimisation complète lancée en arrière-plan."})


@app.post("/api/models/unload")
@admin_required
def api_unload_models():
    embedder.unload()
    reranker.unload()
    return jsonify({"message": "Modèles libérés de la RAM. Ils seront rechargés au prochain besoin."})


@app.get("/api/health")
def api_health():
    return jsonify({"ok": True})


@app.errorhandler(413)
def too_large(_):
    return jsonify({
        "error": f"Bloc trop volumineux. L’application envoie les fichiers par blocs de {config.UPLOAD_CHUNK_MB} Mo."
    }), 413


if __name__ == "__main__":
    app.run(host=config.HOST, port=config.PORT, debug=False, threaded=True)