"""
Thread-safe global state management.

Wraps two module-level dicts (`_active_research`, `_termination_flags`)
with accessor functions protected by a single ``threading.RLock``.
All external code should use the accessor functions instead of touching the
dicts directly.
"""

import threading

from sqlalchemy.orm import Session

# ---------------------------------------------------------------------------
# Internal state — never import these directly from other modules
# ---------------------------------------------------------------------------
_active_research: dict[int, dict] = {}
_termination_flags: dict[int, bool] = {}

# A single lock protects both dicts. They have strongly correlated lifecycles
# (entries are created/destroyed together) and operations are fast dict lookups,
# so a single lock is simpler and eliminates any deadlock risk from lock ordering.
_lock = threading.RLock()


# ===================================================================
# active_research accessors
# ===================================================================


def is_research_active(research_id):
    """Return True if *research_id* is in the active-research dict."""
    with _lock:
        return research_id in _active_research


def get_active_research_ids():
    """Return a list of all active research IDs (snapshot)."""
    with _lock:
        return list(_active_research.keys())


def get_active_research_snapshot(research_id):
    """Return a safe snapshot of an active-research entry, or ``None``.

    The returned dict contains only serialisable fields (``thread`` is
    excluded).  The ``log`` list is shallow-copied — individual entries
    are never mutated after creation, so this is safe.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is None:
            return None
        return {
            "progress": entry.get("progress", 0),
            "status": entry.get("status"),
            "log": list(entry.get("log", [])),
            "settings": dict(s)
            if (s := entry.get("settings")) is not None
            else None,
        }


def get_research_field(research_id, field, default=None):
    """Return a single field from an active-research entry.

    For mutable fields (``list``, ``dict``) a shallow copy is returned so
    callers cannot accidentally mutate shared state.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is None:
            return default
        value = entry.get(field, default)
        # We explicitly check for list/dict rather than using copy.copy()
        # because entries may contain threading.Thread objects —
        # copy.copy(Thread) creates a broken shallow copy sharing the same
        # OS thread ident.  Scalars (int, str, bool) are immutable and safe.
        # For Thread objects, access them via iter_active_research() snapshots.
        if isinstance(value, list):
            return list(value)
        if isinstance(value, dict):
            return dict(value)
        return value


def set_active_research(research_id, data):
    """Insert or replace the active-research entry for *research_id*."""
    with _lock:
        _active_research[research_id] = data


def check_and_start_research(research_id, data) -> bool:
    """Atomically register a research entry iff no live thread exists.

    If ``_active_research[research_id]`` already holds an entry whose
    ``thread`` is alive, returns ``False`` without mutating state and
    without calling ``.start()``. Otherwise, starts ``data['thread']``
    and writes *data* into the active-research dict, then returns ``True``.

    The entire check-and-start is done under the single shared lock, so
    two concurrent callers with the same *research_id* cannot both pass
    the liveness check and end up with two live threads for the same ID.
    """
    thread = data.get("thread") if isinstance(data, dict) else None
    if thread is None:
        raise ValueError("data must contain a 'thread' entry")
    with _lock:
        entry = _active_research.get(research_id)
        if entry is not None:
            existing = entry.get("thread")
            if existing is not None and existing.is_alive():
                return False
        thread.start()
        _active_research[research_id] = data
    return True


def update_active_research(research_id, **fields):
    """Update one or more fields on an existing entry.

    Silently does nothing if *research_id* is not active.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is not None:
            entry.update(fields)


def append_research_log(research_id, log_entry):
    """Append *log_entry* to the ``log`` list for *research_id*.

    Silently does nothing if *research_id* is not active.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is not None:
            entry.setdefault("log", []).append(log_entry)


def update_progress_if_higher(research_id, new_progress):
    """Atomically update progress only if *new_progress* exceeds current.

    Returns the resulting progress value, or ``None`` if the research is
    not active.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is None:
            return None
        current = entry.get("progress", 0)
        if new_progress is not None and new_progress > current:
            entry["progress"] = new_progress
            return new_progress
        return current


def remove_active_research(research_id):
    """Remove *research_id* from the active-research dict (if present)."""
    with _lock:
        _active_research.pop(research_id, None)


def iter_active_research():
    """Yield ``(research_id, snapshot)`` pairs for all active research.

    Each *snapshot* is a shallow copy of the entry dict, safe to read
    outside the lock.
    """
    with _lock:
        items = [
            (
                rid,
                {
                    "progress": entry.get("progress", 0),
                    "status": entry.get("status"),
                    "log": list(entry.get("log", [])),
                    "settings": dict(s)
                    if (s := entry.get("settings")) is not None
                    else None,
                },
            )
            for rid, entry in _active_research.items()
        ]
    for rid, data in items:
        yield rid, data


def get_active_research_count():
    """Return the number of active research entries."""
    with _lock:
        return len(_active_research)


def get_usernames_with_active_research() -> set:
    """Return set of usernames that have research currently running."""
    with _lock:
        return {
            entry.get("settings", {}).get("username")
            for entry in _active_research.values()
            if entry.get("settings", {}).get("username")
        }


# ===================================================================
# termination_flags accessors
# ===================================================================


def is_termination_requested(research_id):
    """Return ``True`` if termination was requested for *research_id*."""
    with _lock:
        return _termination_flags.get(research_id, False)


def set_termination_flag(research_id):
    """Signal that *research_id* should be terminated."""
    with _lock:
        _termination_flags[research_id] = True


def clear_termination_flag(research_id):
    """Remove the termination flag for *research_id* (if present)."""
    with _lock:
        _termination_flags.pop(research_id, None)


# ===================================================================
# Compound / cleanup helpers
# ===================================================================


def is_research_thread_alive(research_id):
    """Return ``True`` if the research thread for *research_id* is alive.

    Returns ``False`` if the research is not active or has no thread.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is None:
            return False
        thread = entry.get("thread")
        return thread is not None and thread.is_alive()


def update_progress_and_check_active(research_id, new_progress):
    """Atomically update progress (if higher) and check if research is active.

    Returns ``(progress_value, is_active)`` where *progress_value* is the
    resulting progress (or ``None`` if not active) and *is_active* indicates
    whether *research_id* is still in the active-research dict.
    """
    with _lock:
        entry = _active_research.get(research_id)
        if entry is None:
            return (None, False)
        current = entry.get("progress", 0)
        if new_progress is not None and new_progress > current:
            entry["progress"] = new_progress
            return (new_progress, True)
        return (current, True)


def cleanup_research(research_id):
    """Remove *research_id* from both active_research and termination_flags
    atomically under the single shared lock.
    """
    with _lock:
        _active_research.pop(research_id, None)
        _termination_flags.pop(research_id, None)


def reclaim_stale_user_active_research(
    db_session: Session, username, *, grace_cutoff_dt=None, logger=None
):
    """Flip ``UserActiveResearch`` rows whose worker thread is dead.

    Shared between ``research_routes.start_research`` (no grace window)
    and ``chat.routes.send_message`` (30-second grace window) — both
    sites historically iterated the same query / status flip / cleanup
    pattern inline; consolidating here makes the difference (grace vs.
    no grace) explicit instead of two near-duplicate copies drifting.

    The helper does NOT commit — callers compose this with surrounding
    DB writes and commit once at the end. Returns ``True`` if any rows
    were reclaimed so the caller can decide whether a commit is needed.

    Args:
        db_session: open SQLAlchemy session against the user's DB.
        username: the per-user DB owner (used to scope the query).
        grace_cutoff_dt: optional ``datetime`` boundary; rows whose
            ``started_at`` is at-or-after this are skipped (avoids
            killing a sibling request's just-spawned thread). ``None``
            disables the grace filter — matches the
            ``research_routes.start_research`` original behaviour.
        logger: optional ``loguru``-style logger for the reclaim
            audit line. Both call sites log at WARNING so operators can
            trace why an active-research cap was released.

    Returns:
        ``True`` if any row was flipped (caller should commit), else
        ``False``.
    """
    # Lazy import to avoid pulling SQLAlchemy at module import time;
    # this helper is only called from the route paths that already have
    # the ORM models in scope.
    from ...constants import ResearchStatus
    from ...database.models import UserActiveResearch

    query = db_session.query(UserActiveResearch).filter(
        UserActiveResearch.username == username,
        UserActiveResearch.status == ResearchStatus.IN_PROGRESS,
    )
    if grace_cutoff_dt is not None:
        query = query.filter(UserActiveResearch.started_at < grace_cutoff_dt)

    reclaimed = False
    for row in query.all():
        if is_research_thread_alive(row.research_id):
            continue
        if logger is not None:
            logger.warning(
                "Reclaiming stale UserActiveResearch {short_id}... "
                "(thread dead) for user {user}",
                short_id=row.research_id[:8],
                user=username,
            )
        row.status = ResearchStatus.FAILED
        cleanup_research(row.research_id)
        reclaimed = True
    return reclaimed
