#!/usr/bin/env python3
"""Flask + MCP retrieval + WaterScope agent UI."""

from __future__ import annotations

import os
from pathlib import Path

from dotenv import load_dotenv

load_dotenv(override=True)


# ---------------------------------------------------------------------------
# Hugging Face Dataset bootstrap
# Download only the search-critical metadata + Qdrant index at startup.
# The large PDFs/Summary/Figures assets are intentionally NOT downloaded here.
# ---------------------------------------------------------------------------
def _bootstrap_search_data() -> None:
    if os.getenv("WATERSCOPE_DATA_BOOTSTRAP", "1").lower() not in {"1", "true", "yes"}:
        return

    from huggingface_hub import HfApi, hf_hub_download
    from huggingface_hub.hf_api import RepoFile
    from time import monotonic

    repo_id = os.getenv(
        "WATERSCOPE_DATASET_REPO",
        "EnvironmentalAI/WaterScopeAI-Data",
    ).strip()
    data_root = Path(
        os.getenv("WATERSCOPE_DATA_ROOT", "/tmp/waterscope-data")
    ).resolve()
    data_root.mkdir(parents=True, exist_ok=True)

    started = monotonic()

    def log(message: str) -> None:
        print(f"[bootstrap pid={os.getpid()} +{monotonic() - started:.1f}s] {message}", flush=True)

    api = HfApi()
    index_prefix = "metadata_search_data/qdrant/"
    log(f"Starting search-data bootstrap for {repo_id}")

    def download(filename: str) -> None:
        log(f"Downloading {filename}")
        try:
            hf_hub_download(
                repo_id=repo_id,
                repo_type="dataset",
                filename=filename,
                local_dir=str(data_root),
            )
        except Exception as exc:
            log(f"Download failed for {filename}: {type(exc).__name__}: {exc}")
            raise
        log(f"Ready: {filename}")

    # Known metadata paths do not require a repository-wide file listing.
    download("metadata/papers.csv")
    download("metadata_search_data/qdrant_manifest.json")

    log(f"Listing only {index_prefix}")
    try:
        index_files = []
        for entry in api.list_repo_tree(
            repo_id=repo_id,
            repo_type="dataset",
            path_in_repo=index_prefix.rstrip("/"),
            recursive=True,
            expand=False,
        ):
            if isinstance(entry, RepoFile) and entry.path.startswith(index_prefix):
                index_files.append(entry.path)
                log(f"Found index file: {entry.path}")
    except Exception as exc:
        log(f"Qdrant listing failed: {type(exc).__name__}: {exc}")
        raise

    if not index_files:
        raise RuntimeError(f"No Qdrant files found in Dataset under {index_prefix}")
    log(f"Found {len(index_files)} Qdrant files; media/PDF/summary folders were not listed")
    for filename in sorted(set(index_files)):
        download(filename)

    paper_csv = data_root / "metadata" / "papers.csv"
    qdrant_path = data_root / "metadata_search_data" / "qdrant"
    qdrant_manifest = data_root / "metadata_search_data" / "qdrant_manifest.json"

    if not paper_csv.is_file():
        raise RuntimeError(f"Dataset bootstrap did not produce papers.csv: {paper_csv}")
    if not qdrant_path.is_dir():
        raise RuntimeError(f"Dataset bootstrap did not produce Qdrant directory: {qdrant_path}")

    # Set these before importing metadata_search.config/search modules.
    os.environ["PAPER_CSV"] = str(paper_csv)
    os.environ["METADATA_QDRANT_PATH"] = str(qdrant_path)
    os.environ["METADATA_QDRANT_MANIFEST"] = str(qdrant_manifest)

    print(f"[bootstrap] PAPER_CSV={paper_csv}")
    print(f"[bootstrap] METADATA_QDRANT_PATH={qdrant_path}")
    print(f"[bootstrap] METADATA_QDRANT_MANIFEST={qdrant_manifest}")
    log("Search-data bootstrap complete")


_bootstrap_search_data()

from flask import Flask, jsonify, request, send_from_directory

# Import config FIRST, immediately after the Dataset bootstrap.
# This prevents another metadata_search import chain from loading config
# before the bootstrap environment variables are resolved.
from metadata_search.config import METADATA_SEARCH_PORT, PAPER_CSV, PAPER_FIGURES_DIR, PROJECT_ROOT

from metadata_search.chat_api import chat_payload, get_papers_payload
from metadata_search.paper_figures import (
    figure_paper_count,
    figures_payload_for_paper,
    figures_dir,
    resolve_figure_asset,
)
from metadata_search.mcp_retrieval import (
    list_saved_papers,
    remove_saved_paper,
    save_paper_record,
    search_papers,
)
from metadata_search.search_service import search_score_mode
from backend.mcp_sync import use_mcp
from backend.model import llm_backend_name

DEFAULT_TOP_K = 10
ALLOWED_TOP_K = {10, 50, 100, 200}

if not Path(PAPER_CSV).is_file():
    raise SystemExit(f"Paper dataset not found: {PAPER_CSV}. Set PAPER_CSV in .env.")

app = Flask(__name__)

# Optional multimodal figure/table analysis test endpoint
from backend.media_analysis.test_routes import media_test_bp
app.register_blueprint(media_test_bp)

# Public media route is gated by WATERSCOPE_ENABLE_PUBLIC_MEDIA.
from backend.media_analysis.public_routes import media_public_bp
app.register_blueprint(media_public_bp)


UI_DIR = PROJECT_ROOT / "github_pages"

def _page(name: str):
    return send_from_directory(UI_DIR, name)


def _parse_top_k(raw: str | None) -> int:
    try:
        value = int(raw or DEFAULT_TOP_K)
    except (TypeError, ValueError):
        return DEFAULT_TOP_K
    return value if value in ALLOWED_TOP_K else DEFAULT_TOP_K


@app.get("/")
def home_page():
    if UI_DIR.is_dir():
        return _page("index.html")
    return jsonify(
        {
            "service": "WaterScopeAI Search System",
            "status": "running",
            "health": "/api/health",
            "search": "/api/search?q=activated+sludge&top_k=10",
        }
    )


@app.get("/demo")
def demo_page():
    return _page("demo.html")


@app.get("/chat")
def chat_page():
    return _page("chat.html")


@app.get("/instructions")
def instructions_page():
    return _page("instructions.html")


@app.get("/agent")
def agent_page():
    return _page("agent.html")


@app.get("/Data/<path:filename>")
def mcqa_data(filename: str):
    return send_from_directory(UI_DIR / "Data", filename)


@app.get("/assets/<path:filename>")
def assets(filename: str):
    response = send_from_directory(UI_DIR, filename)
    if filename.endswith((".js", ".css")):
        response.headers["Cache-Control"] = "no-cache"
    return response


@app.before_request
def handle_api_preflight():
    if request.method == "OPTIONS" and request.path.startswith("/api/"):
        return ("", 204)


def _cors_allow_origin() -> str | None:
    configured = os.getenv("WATERSCOPE_CORS_ORIGIN", "*").strip()
    origin = request.headers.get("Origin")
    if configured in {"*", ""}:
        return origin or "*"
    allowed = {item.strip() for item in configured.split(",") if item.strip()}
    if origin and (origin in allowed or "localhost" in origin or "127.0.0.1" in origin):
        return origin
    return configured if configured in allowed else None


@app.after_request
def cors_headers(response):
    if os.getenv("WATERSCOPE_CORS", "0").lower() in {"1", "true", "yes"}:
        allow_origin = _cors_allow_origin()
        if allow_origin:
            response.headers["Access-Control-Allow-Origin"] = allow_origin
        response.headers["Access-Control-Allow-Headers"] = "Content-Type"
        response.headers["Access-Control-Allow-Methods"] = "GET, POST, DELETE, OPTIONS"
    return response


def _year_from_asset_filename(file_name: str) -> str | None:
    import re

    match = re.search(r"\s-\s(\d{4})\s-\s", Path(file_name).name)
    return match.group(1) if match else None


def _download_pdf_on_demand(file_name: str) -> Path | None:
    """Download one PDF from the HF Dataset and cache it locally."""
    if os.getenv("WATERSCOPE_PDF_ON_DEMAND", "1").lower() not in {"1", "true", "yes"}:
        return None

    year = _year_from_asset_filename(file_name)
    if not year:
        return None

    from huggingface_hub import hf_hub_download

    repo_id = os.getenv(
        "WATERSCOPE_DATASET_REPO",
        "EnvironmentalAI/WaterScopeAI-Data",
    ).strip()
    pdf_prefix = os.getenv("WATERSCOPE_PDF_PREFIX", "PDFs").strip() or "PDFs"
    cache_root = Path(
        os.getenv("WATERSCOPE_PDF_CACHE", "/tmp/waterscope-pdf-cache")
    ).resolve()
    cache_root.mkdir(parents=True, exist_ok=True)

    repo_file = f"{pdf_prefix}/{year}/{Path(file_name).name}"

    try:
        downloaded = hf_hub_download(
            repo_id=repo_id,
            repo_type="dataset",
            filename=repo_file,
            local_dir=str(cache_root),
        )
    except Exception as exc:
        print(
            f"[pdf] download failed for {repo_file}: {type(exc).__name__}: {exc}",
            flush=True,
        )
        return None

    path = Path(downloaded)
    return path if path.is_file() else None


@app.get("/pdfs/<path:filename>")
def paper_pdf_file(filename: str):
    safe_name = Path(filename).name

    # 1) Preserve existing local behavior first.
    root_raw = os.getenv("PAPER_PDF_DIR", "").strip()
    if root_raw:
        root = Path(root_raw).resolve()
        if root.is_dir():
            target = (root / safe_name).resolve()
            if target.parent == root and target.is_file():
                return send_from_directory(root, safe_name, as_attachment=True)

    # 2) Cloud fallback: download only the requested PDF from the HF Dataset.
    downloaded = _download_pdf_on_demand(safe_name)
    if downloaded is None:
        return jsonify({"error": "PDF file not found."}), 404

    return send_from_directory(
        downloaded.parent,
        downloaded.name,
        as_attachment=True,
    )


@app.get("/Media/<path:paper_key>/<filename>")
@app.get("/figures/<path:paper_key>/<filename>")
def paper_figure_file(paper_key: str, filename: str):
    target = resolve_figure_asset(paper_key, filename)
    if target is None:
        return jsonify({"error": "Figure file not found."}), 404

    return send_from_directory(
        target.parent,
        target.name,
    )


@app.get("/api/papers/figures")
def api_paper_figures():
    txt_file = request.args.get("txt_file", "").strip()
    pdf_file = request.args.get("pdf_file", "").strip()
    paper_key = request.args.get("paper_key", "").strip()
    metadata_row_index_raw = request.args.get("metadata_row_index", "").strip()

    metadata_row_index: int | None = None
    if metadata_row_index_raw:
        try:
            metadata_row_index = int(metadata_row_index_raw)
        except ValueError:
            return jsonify({"error": "metadata_row_index must be an integer."}), 400

    if not any([txt_file, pdf_file, paper_key, metadata_row_index is not None]):
        return jsonify({"error": "Provide txt_file, pdf_file, paper_key, or metadata_row_index."}), 400

    payload = figures_payload_for_paper(
        txt_file=txt_file,
        pdf_file=pdf_file,
        paper_key=paper_key,
        metadata_row_index=metadata_row_index,
    )
    if payload is None:
        return jsonify({"error": "No figures found for this paper."}), 404
    return jsonify(payload)


@app.get("/api/search")
def api_search():
    query = request.args.get("q", "").strip()
    if not query:
        return jsonify({"error": "q is required"}), 400
    top_k = _parse_top_k(request.args.get("top_k"))
    try:
        from metadata_search.boolean_query import parse_boolean_query, plan_summary

        results, score_mode = search_papers(query, top_k=top_k)
        parsed = parse_boolean_query(query)
        return jsonify(
            {
                "query": query,
                "search_query": parsed.raw if parsed.raw != query else query,
                "results": results,
                "score_mode": score_mode,
                "query_parsed": plan_summary(parsed),
            }
        )
    except Exception as exc:
        message = str(exc) or repr(exc) or type(exc).__name__
        return jsonify({"error": message}), 500


@app.get("/api/saved-papers")
def api_saved_papers():
    return jsonify({"papers": list_saved_papers()})


@app.post("/api/saved-papers")
def api_create_saved_paper():
    body = request.get_json(force=True, silent=True) or {}
    try:
        return jsonify({"paper": save_paper_record(body)})
    except ValueError as exc:
        return jsonify({"error": str(exc)}), 400


@app.delete("/api/saved-papers/<paper_id>")
def api_delete_saved_paper(paper_id: str):
    if not remove_saved_paper(paper_id):
        return jsonify({"error": "Saved paper not found."}), 404
    return jsonify({"ok": True})


@app.get("/api/health")
def api_health():
    summaries_dir = os.getenv("PAPER_SUMMARIES_DIR", "")
    from metadata_search.search_service import get_rows, get_vector_store

    paper_rows = 0
    vector_points = 0
    vector_error = None
    rows_error = None

    try:
        paper_rows = len(get_rows())
    except Exception as exc:
        rows_error = f"{type(exc).__name__}: {exc}"

    try:
        vector_points = get_vector_store().points_count()
    except Exception as exc:
        vector_error = f"{type(exc).__name__}: {exc}"
        print(f"[health] Vector-store error: {vector_error}", flush=True)

    return jsonify(
        {
            "status": "ok",
            "checks": {
                "paper_csv": Path(PAPER_CSV).is_file(),
                "summaries_dir": bool(summaries_dir and Path(summaries_dir).is_dir()),
                "metadata_vector_index": vector_points > 0,
                "paper_figures_dir": PAPER_FIGURES_DIR.is_dir(),
            },
            "paper_count": paper_rows,
            "metadata_vector_points": vector_points,
            "figure_paper_count": figure_paper_count(),
            "llm_backend": llm_backend_name(),
            "retrieval": "mcp" if use_mcp() else "direct",
            "search_score_mode": search_score_mode(),
            "diagnostics": {
                "paper_csv_path": str(PAPER_CSV),
                "qdrant_path": os.getenv("METADATA_QDRANT_PATH", ""),
                "qdrant_manifest": os.getenv("METADATA_QDRANT_MANIFEST", ""),
                "rows_error": rows_error,
                "vector_error": vector_error,
            },
        }
    )


@app.get("/api/papers")
def api_papers():
    try:
        return jsonify(get_papers_payload())
    except FileNotFoundError as exc:
        return jsonify({"error": str(exc)}), 503
    except Exception as exc:
        return jsonify({"error": f"Could not load papers: {exc}"}), 500


@app.post("/api/chat")
def api_chat():
    body = request.get_json(force=True, silent=True) or {}
    try:
        return jsonify(chat_payload(body))
    except ValueError as exc:
        message = str(exc)
        if "RUNTIME_ERROR" in message or "invalid state" in message.lower():
            return jsonify(
                {
                    "error": (
                        "WaterScope HF Space is unavailable right now "
                        f"(EnvironmentalAI/WaterScopeAI: {message}). "
                        "Retrieval succeeded, but the answer model could not run."
                    )
                }
            ), 503
        return jsonify({"error": message}), 400
    except Exception as exc:
        return jsonify(
            {
                "error": (
                    "Could not reach the WaterScope HF Space model. "
                    f"Check network and EnvironmentalAI/WaterScopeAI ({type(exc).__name__}: {exc})"
                )
            }
        ), 503


def main() -> None:
    from metadata_search.search_service import auto_rebuild_on_startup

    auto_rebuild_on_startup()
    if use_mcp():
        from backend.mcp_sync import ensure_mcp_connected

        ensure_mcp_connected()
    debug = os.getenv("FLASK_DEBUG", "0").lower() in {"1", "true", "yes"}
    app.run(host=os.getenv("HOST", "0.0.0.0"), debug=debug, use_reloader=False, port=METADATA_SEARCH_PORT)


if __name__ == "__main__":
    main()