"""Session-level data fetching.

All functions take a psycopg.Connection as the first arg and return DataFrames.
"""

from datetime import datetime

import pandas as pd
import psycopg

from .connection import query_df


# ---------------------------------------------------------------------------
# Full-table / overview
# ---------------------------------------------------------------------------

def list_all(conn: psycopg.Connection, limit: int = 200) -> pd.DataFrame:
    """All sessions ordered by last_updated desc."""
    return query_df(conn, """
        SELECT * FROM sessions ORDER BY last_updated DESC LIMIT %s
    """, (limit,))


def list_users(conn: psycopg.Connection) -> pd.DataFrame:
    """Distinct users with session count and last seen."""
    return query_df(conn, """
        SELECT   user_email, user_name, org,
                 COUNT(*)           AS session_count,
                 MAX(last_updated)  AS last_active
        FROM     sessions
        GROUP BY user_email, user_name, org
        ORDER BY last_active DESC
    """)


def list_orgs(conn: psycopg.Connection) -> pd.DataFrame:
    """Distinct orgs with user count and session count."""
    return query_df(conn, """
        SELECT   org,
                 COUNT(DISTINCT user_email) AS user_count,
                 COUNT(*)                   AS session_count,
                 MAX(last_updated)          AS last_active
        FROM     sessions
        GROUP BY org
        ORDER BY session_count DESC
    """)


def list_repos(conn: psycopg.Connection) -> pd.DataFrame:
    """Distinct repos with session count."""
    return query_df(conn, """
        SELECT   repo_name, org,
                 COUNT(*)          AS session_count,
                 MAX(last_updated) AS last_active
        FROM     sessions
        WHERE    repo_name IS NOT NULL
        GROUP BY repo_name, org
        ORDER BY session_count DESC
    """)


# ---------------------------------------------------------------------------
# Filtered fetchers
# ---------------------------------------------------------------------------

def by_user(
    conn: psycopg.Connection,
    email: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 500,
) -> pd.DataFrame:
    """Sessions for a single user, optionally within a date range."""
    clauses = ["user_email = %s"]
    params: list = [email]
    _add_time_range(clauses, params, since, until, col="first_seen")
    sql = f"""
        SELECT * FROM sessions
        WHERE  {' AND '.join(clauses)}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_users(
    conn: psycopg.Connection,
    emails: list[str],
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 1000,
) -> pd.DataFrame:
    """Sessions for multiple users."""
    clauses = ["user_email = ANY(%s)"]
    params: list = [emails]
    _add_time_range(clauses, params, since, until, col="first_seen")
    sql = f"""
        SELECT * FROM sessions
        WHERE  {' AND '.join(clauses)}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_org(
    conn: psycopg.Connection,
    org: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 1000,
) -> pd.DataFrame:
    """Sessions for an org."""
    clauses = ["org = %s"]
    params: list = [org]
    _add_time_range(clauses, params, since, until, col="first_seen")
    sql = f"""
        SELECT * FROM sessions
        WHERE  {' AND '.join(clauses)}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_repo(
    conn: psycopg.Connection,
    repo_name: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 500,
) -> pd.DataFrame:
    """Sessions for a specific repo."""
    clauses = ["repo_name = %s"]
    params: list = [repo_name]
    _add_time_range(clauses, params, since, until, col="first_seen")
    sql = f"""
        SELECT * FROM sessions
        WHERE  {' AND '.join(clauses)}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_device(
    conn: psycopg.Connection,
    device_id: str,
    *,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 500,
) -> pd.DataFrame:
    """Sessions from a specific device."""
    clauses = ["device_id = %s"]
    params: list = [device_id]
    _add_time_range(clauses, params, since, until, col="first_seen")
    sql = f"""
        SELECT * FROM sessions
        WHERE  {' AND '.join(clauses)}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


def by_id(conn: psycopg.Connection, session_id: str) -> pd.DataFrame:
    """Single session by ID."""
    return query_df(conn, "SELECT * FROM sessions WHERE id = %s", (session_id,))


def by_ids(conn: psycopg.Connection, session_ids: list[str]) -> pd.DataFrame:
    """Multiple sessions by IDs."""
    return query_df(conn, "SELECT * FROM sessions WHERE id = ANY(%s)", (session_ids,))


def search(
    conn: psycopg.Connection,
    *,
    email: str | None = None,
    org: str | None = None,
    repo_name: str | None = None,
    source: str | None = None,
    branch: str | None = None,
    device_id: str | None = None,
    model: str | None = None,
    since: datetime | None = None,
    until: datetime | None = None,
    limit: int = 500,
) -> pd.DataFrame:
    """Flexible multi-filter session search. All filters are AND-ed."""
    clauses: list[str] = []
    params: list = []
    _add_eq(clauses, params, "user_email", email)
    _add_eq(clauses, params, "org", org)
    _add_eq(clauses, params, "repo_name", repo_name)
    _add_eq(clauses, params, "source", source)
    _add_eq(clauses, params, "git_branch", branch)
    _add_eq(clauses, params, "device_id", device_id)
    _add_eq(clauses, params, "model", model)
    _add_time_range(clauses, params, since, until, col="first_seen")
    where = f"WHERE {' AND '.join(clauses)}" if clauses else ""
    sql = f"""
        SELECT * FROM sessions {where}
        ORDER  BY first_seen DESC
        LIMIT  %s
    """
    params.append(limit)
    return query_df(conn, sql, tuple(params))


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _add_eq(clauses: list, params: list, col: str, val) -> None:
    if val is not None:
        clauses.append(f"{col} = %s")
        params.append(val)


def _add_time_range(
    clauses: list, params: list,
    since: datetime | None, until: datetime | None,
    col: str,
) -> None:
    if since is not None:
        clauses.append(f"{col} >= %s")
        params.append(since)
    if until is not None:
        clauses.append(f"{col} < %s")
        params.append(until)
