import json
from contextlib import nullcontext
from pathlib import Path
from typing import List
from typing import Optional

import click
from pydantic import Field
from requests import HTTPError
from rich.panel import Panel
from rich.progress import track

from _incydr_cli import console
from _incydr_cli import get_user_project_path
from _incydr_cli import logging_options
from _incydr_cli import render
from _incydr_cli.cmds.options.output_options import columns_option
from _incydr_cli.cmds.options.output_options import input_format_option
from _incydr_cli.cmds.options.output_options import output_options
from _incydr_cli.cmds.options.output_options import single_format_option
from _incydr_cli.cmds.options.output_options import SingleFormat
from _incydr_cli.cmds.options.output_options import table_format_option
from _incydr_cli.cmds.options.output_options import TableFormat
from _incydr_cli.cmds.options.utils import checkpoint_option
from _incydr_cli.cmds.utils import warn_interrupt
from _incydr_cli.core import IncydrCommand
from _incydr_cli.core import IncydrGroup
from _incydr_cli.cursor import CursorStore
from _incydr_cli.logger import get_server_logger
from _incydr_sdk.core.client import Client
from _incydr_sdk.core.models import CSVModel
from _incydr_sdk.core.models import Model
from _incydr_sdk.enums.sessions import ContentInspectionStatuses
from _incydr_sdk.enums.sessions import SessionSeverities
from _incydr_sdk.enums.sessions import SessionStates
from _incydr_sdk.file_events.models.event import FileEventV2
from _incydr_sdk.file_events.models.response import FileEventsPage
from _incydr_sdk.sessions.models.response import Session
from _incydr_sdk.utils import model_as_card

SEVERITY_MAP = {"NO_RISK": 0, "LOW": 1, "MODERATE": 2, "HIGH": 3, "CRITICAL": 4}


@click.group(cls=IncydrGroup)
@logging_options
def sessions():
    """
    View and manage Incydr Sessions, including Alerts.

    Sessions group all Incydr activity for an individual actor within a specific time frame.
    Sessions may include file events, alerts, Instructor lessons, and more. A single session can contain one or more alerts.

    Alerts are triggered when user file activity matches criteria defined in your alert rules.
    Alerts are included within sessions to group related activity and provide additional context.
    """


@sessions.command(cls=IncydrCommand)
@checkpoint_option
@click.option(
    "--actor-id", default=None, help="Limit search to sessions generated by this actor."
)
@click.option(
    "--start",
    default=None,
    help="Limit search to sessions beginning on or after this date and time. "
    "Accepts a date/time in yyyy-MM-dd (UTC) or yyyy-MM-dd HH:MM:SS (UTC+24-hr time) format.",
)
@click.option(
    "--end",
    default=None,
    help="Limit search to sessions beginning before this date and time. "
    "Accepts a date/time in yyyy-MM-dd (UTC) or yyyy-MM-dd HH:MM:SS (UTC+24-hr time) format.",
)
@click.option(
    "--no-alerts",
    is_flag=True,
    default=False,
    help="Limit search to sessions that do NOT have alerts associated with them.",
)
@click.option(
    "--risk-indicators",
    default=None,
    help="A CSV list of risk indicator IDs. Limit search to sessions that include these risk indicators.",
)
@click.option(
    "--state",
    type=click.Choice([e.value for e in SessionStates]),
    default=[],
    multiple=True,
    help="Limit search to sessions with this state. Can be specified multiple times to include multiple values. ",
)
@click.option(
    "--severity",
    type=click.Choice([e.value for e in SessionSeverities]),
    default=[],
    multiple=True,
    help="Specify one or more severity levels, can be specified multiple times to include multiple values. "
    "Limit search to sessions that have the matching severity value(s). ",
)
@click.option(
    "--rule-id",
    default=[],
    multiple=True,
    help="A rule ID. Limit search to sessions associated with this rule. Can be specified multiple times to include multiple values. ",
)
@click.option(
    "--watchlist-id",
    default=[],
    multiple=True,
    help="A watchlist ID. Limit search to sessions associated with this watchlist. Can be specified multiple times to include multiple values. ",
)
@click.option(
    "--content-inspection-status",
    type=click.Choice([e.value for e in ContentInspectionStatuses]),
    default=None,
    help="Limit search to sessions with the given content inspection status.",
)
@table_format_option
@columns_option
@logging_options
def search(
    actor_id: Optional[str] = None,
    start: Optional[str] = None,
    end: Optional[str] = None,
    no_alerts: bool = False,
    risk_indicators: Optional[str] = None,
    state: Optional[List[str]] = None,
    severity: Optional[List[str]] = None,
    rule_id: Optional[List[str]] = None,
    watchlist_id: Optional[List[str]] = None,
    content_inspection_status: Optional[str] = None,
    output: Optional[str] = None,
    certs: Optional[str] = None,
    ignore_cert_validation: Optional[bool] = None,
    checkpoint_name: Optional[str] = None,
    format_: Optional[TableFormat] = None,
    columns: Optional[str] = None,
):
    """
    Search sessions.

    Defaults to only include sessions that have alerts associated with them.
    Use the --no-alerts option to view sessions without any alerts.
    """
    client = Client()
    cursor = _get_cursor_store(client.settings.api_client_id)

    if output:
        format_ = TableFormat.json_lines

    # Use stored checkpoint timestamp for start filter if applicable
    if checkpoint_name:
        checkpoint = cursor.get(checkpoint_name)
        if checkpoint:
            start = int(checkpoint)

    severity_values = []
    for value in severity:
        severity_values.append(SEVERITY_MAP[value])

    sessions_gen = client.sessions.v1.iter_all(
        actor_id=actor_id,
        start_time=start,
        end_time=end,
        has_alerts=not no_alerts,
        risk_indicators=risk_indicators.split(",") if risk_indicators else None,
        states=[*state] if state else None,
        severities=severity_values,
        rule_ids=[*rule_id] if rule_id else None,
        watchlist_ids=[*watchlist_id] if watchlist_id else None,
        content_inspection_status=content_inspection_status,
    )
    if checkpoint_name:
        sessions_gen = _update_checkpoint(cursor, checkpoint_name, sessions_gen)

    with warn_interrupt() if checkpoint_name else nullcontext():
        if output:
            logger = get_server_logger(output, certs, ignore_cert_validation)
            for session in sessions_gen:
                logger.info(session.json())
            return

        if format_ == TableFormat.table:
            columns = columns or [
                "session_id",
                "begin_time",
                "end_time",
                "context_summary",
                "exfiltration_summary",
                "actor_id",
                "low_events",
                "moderate_events",
                "high_events",
                "states",
                "notes",
            ]
            render.table(Session, sessions_gen, columns=columns, flat=False)
        elif format_ == TableFormat.csv:
            render.csv(Session, sessions_gen, columns=columns, flat=True)
        else:
            printed = False
            for session in sessions_gen:
                printed = True
                if format_ == TableFormat.json_pretty:
                    console.print_json(session.json())
                else:
                    click.echo(session.json())
            if not printed:
                console.print("No results found.")


@sessions.command(cls=IncydrCommand)
@click.argument("session-id")
@logging_options
@single_format_option
def show(session_id: str, format_: SingleFormat):
    """
    Show the details of a single session.
    """
    client = Client()
    session = client.sessions.v1.get_session_details(session_id)

    if format_ == SingleFormat.rich:
        console.print(Panel.fit(model_as_card(session)))
    elif format_ == SingleFormat.json_pretty:
        console.print_json(session.json())
    else:
        click.echo(session.json())


@sessions.command(cls=IncydrCommand)
@logging_options
@columns_option
@table_format_option
@output_options
@click.argument("session-id")
def show_events(
    session_id: str,
    columns: Optional[str],
    output: Optional[str],
    certs: Optional[str],
    ignore_cert_validation: Optional[bool],
    format_: SingleFormat,
):
    """
    Show the details of a single session.
    """
    client = Client()
    page = client.sessions.v1.get_session_events(session_id)

    def event_iterator(event_page: FileEventsPage):
        yield from event_page.file_events

    events = event_iterator(page)

    if output:
        logger = get_server_logger(output, certs, ignore_cert_validation)
        for event_ in events:
            logger.info(event_.json())
        return

    if format_ == TableFormat.csv:
        render.csv(FileEventV2, events, columns=columns, flat=True)
    elif format_ == TableFormat.table:
        render.table(FileEventV2, events, columns=columns, flat=False)
    else:
        printed = False
        for event in events:
            printed = True
            if format_ == TableFormat.json_pretty:
                console.print_json(data=event)
            else:
                click.echo(json.dumps(event))
        if not printed:
            console.print("No results found.")


@sessions.command(cls=IncydrCommand)
@click.argument("session-id")
@click.option(
    "--state",
    "-s",
    type=click.Choice([e.value for e in SessionStates]),
    help="Update a session's state.",
)
@click.option(
    "--note",
    help="Add a note to the session. Takes the notes content with a 2000 character max.",
)
@logging_options
def update(session_id: str, state: str = None, note: str = None):
    """
    Update the state of and/or the note attached to the session.
    """
    client = Client()
    if state:
        client.sessions.v1.update_state_by_id(session_id, state)
        console.print(f"Successfully updated session {session_id} to {state}!")
    if note:
        client.sessions.v1.add_note(session_id, note)
        console.print(f"Successfully added note to session {session_id}!")


@sessions.command(cls=IncydrCommand)
@click.argument("file", type=click.File())
@click.option(
    "--state",
    type=click.Choice([e.value for e in SessionStates]),
    help="Override CSV/JSON input's `state` value with this value.",
)
@click.option("--note", help="Override CSV/JSON input's `note` value with this value.")
@input_format_option
@logging_options
def bulk_update_state(
    file: Path,
    state: Optional[str] = None,
    note: Optional[str] = None,
    format_: str = None,
):
    """
    Bulk update the state of multiple sessions. Optionally attach a note.

    NEW_STATE specifies the new state to which sessions will be updated.
    Must be one of the following: 'OPEN', 'IN_PROGRESS', 'CLOSED', 'CLOSED_TP', 'CLOSED_FP', 'OPEN_NEW_DATA'

    Takes a single arg `FILE` which specifies the path to the file (use "-" to read from stdin).
    File format can either be CSV or [JSON Lines format](https://jsonlines.org) (Default is CSV).

    The --state and --note options to this command will override respective columns/keys in the CSV/JSON input.

    Accepts the following columns:

    * `session_id` (required)
    * `state` (required if --state option is not provided)
    * `note` (optional)

    If --state is not provided, the CSV/JSON input _must_ have a `state` column/key for each row/object.

    """
    # if --state is provided, we want that column/key to be optional on input data, otherwise required
    state_type = Optional[str] if state else str

    # Validate input models

    class SessionCSV(CSVModel):
        session_id: str = Field(csv_aliases=["session_id", "sessionId"])
        state: state_type = Field(csv_aliases=["state"])
        note: Optional[str]

    class SessionJSON(Model):
        session_id: str = Field(alias="sessionId")
        state: state_type
        note: Optional[str]

    if format_ == "csv":
        models = SessionCSV.parse_csv(file)
    else:
        models = SessionJSON.parse_json_lines(file)

    # Process input

    client = Client()

    # optimization: bucketize update state calls in batches of 100

    try:
        for row in track(
            models,
            description="Updating sessions...",
            transient=True,
        ):
            client.sessions.v1.update_state_by_id(row.session_id, state or row.state)
            if note or row.note:
                client.sessions.v1.add_note(row.session_id, note or row.note)
    except HTTPError as err:
        console.print(f"[red]Error:[/red] {err.response.text}")


def _get_cursor_store(api_key):
    """
    Get cursor store for sessions search checkpoints.
    """
    dir_path = get_user_project_path(
        "checkpoints",
        api_key,
        "session_checkpoints",
    )
    return CursorStore(dir_path, "sessions")


def _update_checkpoint(cursor, checkpoint_name, sessions_gen):
    """
    De-duplicates results across checkpointed runs.

    Since using the timestamp of the last result
    processed as the `--start` time of the next run causes the last event to show up again in the
    next results, store the last session IDs in the cursor to
    filter out on the next run.

    It's also possible that two sessions have the exact same timestamp, so
    `checkpoint_sessions` needs to be a list of session IDs so we can filter out everything that's actually
    been processed.
    """
    checkpoint_sessions = cursor.get_items(checkpoint_name)
    new_timestamp = None
    new_sessions = []
    for session in sessions_gen:
        session_id = session.session_id
        if session_id not in checkpoint_sessions:
            if not new_timestamp or session.end_time > new_timestamp:
                new_timestamp = session.end_time
                new_sessions.clear()
            new_sessions.append(session_id)
            yield session
            cursor.replace(checkpoint_name, new_timestamp)
            cursor.replace_items(checkpoint_name, new_sessions)
