"""Correlate direct RDP logoff (Security 4634/4647) with workstation rdp.login.success."""

from __future__ import annotations

from sqlalchemy import select
from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified

from app.models import Event
from app.services.host_sessions import (
    SESSION_TERMINATED_AT_KEY,
    _details_dict,
    _event_login_user,
)
from app.services.rdg_workstation_session import (
    WORKSTATION_LOGIN_TYPE,
    event_closed_by_rdg,
    users_match_rdg,
)

SESSION_CLOSED_BY_LOGOFF_AT_KEY = "session_closed_by_logoff_at"
SESSION_CLOSED_BY_LOGOFF_EVENT_ID_KEY = "session_closed_by_logoff_event_id"

LOGOFF_EVENT_TYPE = "rdp.session.logoff"
LOGOFF_EVENT_TYPES = frozenset({LOGOFF_EVENT_TYPE})


def event_closed_by_logoff(event: Event) -> bool:
    details = _details_dict(event)
    at = details.get(SESSION_CLOSED_BY_LOGOFF_AT_KEY)
    return at is not None and str(at).strip() != ""


def mark_login_closed_by_logoff(login_event: Event, *, logoff_event: Event) -> None:
    details = dict(_details_dict(login_event))
    details[SESSION_CLOSED_BY_LOGOFF_AT_KEY] = logoff_event.occurred_at.isoformat()
    details[SESSION_CLOSED_BY_LOGOFF_EVENT_ID_KEY] = logoff_event.id
    login_event.details = details
    flag_modified(login_event, "details")


def _login_already_closed(login_event: Event) -> bool:
    details = _details_dict(login_event)
    if details.get("session_terminated") is True:
        return True
    at = details.get(SESSION_TERMINATED_AT_KEY)
    if at is not None and str(at).strip() != "":
        return True
    if event_closed_by_rdg(login_event):
        return True
    return event_closed_by_logoff(login_event)


def find_workstation_login_for_logoff(db: Session, logoff_event: Event) -> Event | None:
    if logoff_event.type not in LOGOFF_EVENT_TYPES:
        return None

    logoff_user = _event_login_user(logoff_event)
    if not logoff_user:
        return None
    logoff_at = logoff_event.occurred_at
    host_id = logoff_event.host_id
    if host_id is None:
        return None

    candidates = db.scalars(
        select(Event)
        .where(
            Event.host_id == host_id,
            Event.type == WORKSTATION_LOGIN_TYPE,
            Event.occurred_at <= logoff_at,
            Event.id != logoff_event.id,
        )
        .order_by(Event.occurred_at.desc())
    ).all()

    logoff_details = _details_dict(logoff_event)
    logoff_ip = str(logoff_details.get("ip_address") or "").strip()

    for login in candidates:
        if _login_already_closed(login):
            continue
        login_user = _event_login_user(login)
        if not users_match_rdg(login_user, logoff_user):
            continue
        if logoff_ip and logoff_ip not in ("", "-"):
            login_ip = str(_details_dict(login).get("ip_address") or "").strip()
            if login_ip and login_ip not in ("", "-") and login_ip != logoff_ip:
                continue
        return login
    return None


def find_logoff_after_workstation_login(db: Session, login_event: Event) -> Event | None:
    """Runtime lookup for historical logins without persisted close flag."""
    if login_event.type != WORKSTATION_LOGIN_TYPE:
        return None
    if _login_already_closed(login_event):
        return None

    host_id = login_event.host_id
    if host_id is None:
        return None

    login_user = _event_login_user(login_event)
    if not login_user:
        return None
    login_at = login_event.occurred_at
    login_ip = str(_details_dict(login_event).get("ip_address") or "").strip()

    candidates = db.scalars(
        select(Event)
        .where(
            Event.host_id == host_id,
            Event.type == LOGOFF_EVENT_TYPE,
            Event.occurred_at >= login_at,
            Event.id != login_event.id,
        )
        .order_by(Event.occurred_at.asc())
    ).all()

    for logoff in candidates:
        if not users_match_rdg(_event_login_user(logoff), login_user):
            continue
        logoff_ip = str(_details_dict(logoff).get("ip_address") or "").strip()
        if login_ip and login_ip not in ("", "-") and logoff_ip and logoff_ip not in ("", "-"):
            if login_ip != logoff_ip:
                continue
        return logoff
    return None


def resolve_workstation_login_closed_by_logoff(db: Session, login_event: Event) -> bool:
    if event_closed_by_logoff(login_event):
        return True
    return find_logoff_after_workstation_login(db, login_event) is not None


def close_workstation_session_for_rdp_logoff(db: Session, logoff_event: Event) -> Event | None:
    """On direct RDP logoff, mark matching open workstation login as session-closed."""
    if logoff_event.type not in LOGOFF_EVENT_TYPES:
        return None

    login = find_workstation_login_for_logoff(db, logoff_event)
    if login is None:
        return None
    mark_login_closed_by_logoff(login, logoff_event=logoff_event)
    return login