from __future__ import annotations
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional

from bson import ObjectId
from pymongo import DESCENDING

from connections import get_runs_collection

log = logging.getLogger(__name__)


# ------------------- Index Setup -------------------
def _ensure_runs_indexes(col):
    try:
        col.create_index([("updated_at", DESCENDING)])
        col.create_index([("flow", 1), ("updated_at", DESCENDING)])
        col.create_index([("created_by", 1), ("updated_at", DESCENDING)])
    except Exception:
        pass


_rc = get_runs_collection()
if _rc is not None:
    _ensure_runs_indexes(_rc)


# ------------------- Runs -------------------
def upsert_run(
    run_id: str,
    *,
    flow: str,
    state_json: str,
    stages: List[Dict[str, Any]],
    urls: Dict[str, Any],
    created_by: Optional[str] = None,
) -> bool:
    """Insert or update a run. Returns False when there is no database."""
    col = get_runs_collection()
    if col is None:
        log.error("run %s not saved: no database", run_id)
        return False
    now = datetime.utcnow()
    set_fields: Dict[str, Any] = {
        "flow": flow,
        "state": state_json,
        "stages": stages,
        "urls": urls,
        "updated_at": now,
    }
    try:
        col.update_one(
            {"_id": ObjectId(run_id)},
            {"$set": set_fields, "$setOnInsert": {"created_by": created_by, "created_at": now}},
            upsert=True,
        )
        return True
    except Exception as e:
        log.error(f"upsert_run failed: {e}")
        return False


def find_run(run_id: str) -> Optional[Dict[str, Any]]:
    col = get_runs_collection()
    if col is None:
        return None
    return col.find_one({"_id": ObjectId(run_id)})


def find_runs(
    *,
    flow: Optional[str] = None,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
    limit: int = 200,
    created_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
    """Runs newest first, optionally filtered by flow ("Research"/"Upload"), user and date."""
    col = get_runs_collection()
    if col is None:
        return []
    q: Dict[str, Any] = {"updated_at": {"$exists": True}}  # skip docs from before this shape
    if flow:
        q["flow"] = flow
    if created_by:
        q["created_by"] = created_by
    if start_date or end_date:
        rng: Dict[str, Any] = {}
        if start_date: rng["$gte"] = start_date
        if end_date:   rng["$lt"]  = end_date
        q["updated_at"] = rng
    cur = col.find(q).sort("updated_at", DESCENDING).limit(max(1, int(limit)))
    out: List[Dict[str, Any]] = []
    for d in cur:
        d["_id"] = str(d.get("_id"))
        out.append(d)
    return out


def delete_run(run_id: str) -> None:
    col = get_runs_collection()
    if col is None:
        return
    try:
        col.delete_one({"_id": ObjectId(run_id)})
    except Exception as e:
        log.error(f"delete_run failed: {e}")