# -*- coding: utf-8 -*-
"""Collects Timesketch events."""

import datetime
import tempfile
from typing import List, Callable

import pandas as pd
from timesketch_api_client import client, search, sketch

from dftimewolf.lib import module, timesketch_utils
from dftimewolf.lib import cache
from dftimewolf.lib import telemetry
from dftimewolf.lib.containers import manager as container_manager
from dftimewolf.lib.containers import containers
from dftimewolf.lib.modules import manager as modules_manager

_VALID_OUTPUT_FORMATS = frozenset(["csv", "json", "jsonl", "pandas"])


def GenerateTimerangeQuery(
  start_datetime: datetime.datetime | None,
  end_datetime: datetime.datetime | None,
) -> str:
  """Returns a Timesketch query string for the given time range.

  Args:
    start_datetime: The start datetime.
    end_datetime: The end datetime.

  Returns:
    A Timesketch query string for the given time range.
  """
  from_query = "*"
  to_query = "*"
  if start_datetime:
    from_query = start_datetime.strftime("%Y-%m-%dT%H:%M:%S")
  if end_datetime:
    to_query = end_datetime.strftime("%Y-%m-%dT%H:%M:%S")
  return f"datetime:[{from_query} TO {to_query}]"


class TimesketchSearchEventCollector(module.BaseModule):
  """Collector for Timesketch events.

  Collects Timesketch events using given search.

  Attributes:
    query_string: the Timesketch query string.
    start_datetime: the start datetime.
    end_datetime: the end datetime.
    indices: the list of Timesketch indices.
    labels: the Timesketch labels.
    output_format: the output format.
    return_fields: the comma-separated Timesketch return fields.
    search_name: an optional name for the search.
    search_description: an optional description for the search.
    include_internal_columns: show Timesketch internal columns.
    sketch_id: the Timesketch sketch ID.
    sketch: the Timesketch sketch.
  """

  def __init__(self,
               name: str,
               container_manager_: container_manager.ContainerManager,
               cache_: cache.DFTWCache,
               telemetry_: telemetry.BaseTelemetry,
               publish_message_callback: Callable[[str, str, bool], None]):
    """Initialise a TimesketchSearchEventCollector."""
    super().__init__(name=name,
                     cache_=cache_,
                     container_manager_=container_manager_,
                     telemetry_=telemetry_,
                     publish_message_callback=publish_message_callback)
    self.query_string: str = ""
    self.start_datetime: datetime.datetime | None = None
    self.end_datetime: datetime.datetime | None = None
    self.indices: List[int] = []
    self.labels: List[str] = []
    self.output_format: str = ""
    self.return_fields: str = ""
    self.search_name: str = ""
    self.search_description: str = ""
    self.include_internal_columns: bool = False
    self.sketch_id: int = 0
    self.sketch: sketch.Sketch | None = None
    self.timesketch_api_client: client.TimesketchApi | None = None

  # pylint: disable=arguments-differ,too-many-arguments
  def SetUp(
    self,
    sketch_id: str | None = None,
    query_string: str = "*",
    start_datetime: datetime.datetime | None = None,
    end_datetime: datetime.datetime | None = None,
    indices: str | None = None,
    labels: str | None = None,
    output_format: str = "pandas",
    return_fields: str = "*",
    search_name: str | None = None,
    search_description: str | None = None,
    include_internal_columns: bool = False,
    token_password: str = "",
    endpoint: str | None = None,
    username: str | None = None,
    password: str | None = None,
  ) -> None:
    """Sets up the TimesketchSearchEventCollector.

    Args:
      sketch_id: the Timesketch sketch ID.  Defaults to the value
          in the ticket attribute container.
      query_string: the query string.  Defaults to '*' (all events).
      start_datetime: the start datetime.
      end_datetime: the end datetime.
      indices: the comma-separated Timesketch indices.
      labels: the comma-separated Timesketch event labels.
      output_format: the output format.  Defaults to 'pandas'.
      return_fields: the return fields.  Defaults to '*'.
      search_name: an optional name for the search.
      search_description: an optional description for the search.
      include_internal_columns: include Timesketch internal columns.  Defaults
          to False.
      token_password: optional password used to decrypt the
          Timesketch credential storage. Defaults to an empty string since
          the upstream library expects a string value. An empty string means
          a password will be generated by the upstream library.
      endpoint: Timesketch server URL (e.g. http://localhost:5000/).
          Optional when token_password is provided.
      username: Timesketch username. Optional when token_password is provided.
      password: Timesketch password. Optional when token_password is provided.
    """
    self.timesketch_api_client = self._GetAPIClient(
      token_password, endpoint, username, password
    )

    self.sketch_id = int(sketch_id) if sketch_id else 0
    self.FindSketch()

    if output_format not in _VALID_OUTPUT_FORMATS:
      self.ModuleError(
        f"Output format not one of {','.join(_VALID_OUTPUT_FORMATS)}",
        critical=True,
      )
    self.start_datetime = start_datetime
    self.end_datetime = end_datetime
    self.query_string = query_string
    self.return_fields = return_fields
    self.output_format = output_format
    self.include_internal_columns = include_internal_columns

    if labels:
      self.labels = [label.strip() for label in labels.split(",")]

    if indices:
      self.indices = [int(index) for index in indices.split(",")]

    if search_name:
      self.search_name = search_name

    if search_description:
      self.search_description = search_description

  def _GetAPIClient(
    self,
    token_password: str | None = None,
    endpoint: str | None = None,
    username: str | None = None,
    password: str | None = None,
  ) -> client.TimesketchApi:
    """Gets the Timesketch API client."""
    if endpoint and username and password:
      timesketch_api = client.TimesketchApi(endpoint, username, password)
    elif token_password:
      timesketch_api = timesketch_utils.GetApiClient(
        self._cache, token_password=token_password
      )
    else:
      timesketch_api = timesketch_utils.GetApiClient(self._cache)

    if not timesketch_api:
      self.ModuleError(
        "Unable to get a Timesketch API client, try deleting the files "
        "~/.timesketchrc and ~/.timesketch.token",
        critical=True,
      )
    if not timesketch_api.session:
      self.ModuleError("Could not connect to Timesketch server.", critical=True)
    return timesketch_api

  def _GetSketch(self, sketch_id: int) -> sketch.Sketch:
    """Gets the Timesketch sketch.

    Args:
      sketch_id: The Timesketch sketch ID.

    Returns:
      The Timesketch sketch.
    """
    assert self.timesketch_api_client
    sketch_obj = self.timesketch_api_client.get_sketch(sketch_id)
    if not sketch_obj:
      self.ModuleError(f"Could not get sketch {sketch_id}", critical=True)
    self._cache.AddToCache("timesketch_sketch", sketch_obj)
    self.logger.warning(f"Adding sketch {sketch_obj} to cache")
    return sketch_obj

  def _GetSearchResults(
    self,
    selected_sketch: sketch.Sketch,
    query_string: str,
    return_fields: str,
    start_datetime: datetime.datetime | None = None,
    end_datetime: datetime.datetime | None = None,
    labels: list[str] | None = None,
    indices: list[int] | None = None,
  ) -> pd.DataFrame:
    """Get the Timesketch search results.

    Args:
      sketch: the Timesketch sketch.
      query_string: the query string.
      return_fields: fields of the sketch to return in the results.
      start_datetime: Optional start datetime filter.
      end_datetime: Optional end datetime filter.
      labels: Filter labels. Can also filter on special labels 'star' and
        'comment'.
      indices: Optional indices to filter on.

    Returns:
      the results in a Pandas dataframe.
    """
    search_obj = search.Search(selected_sketch)
    search_obj.return_fields = return_fields
    if indices:
      search_obj.indices = indices

    date_query = GenerateTimerangeQuery(start_datetime, end_datetime)
    search_obj.query_string = f"({query_string}) AND ({date_query})"

    for label in labels or []:
      label_chip = search.LabelChip()
      if label == "star":
        label_chip.use_star_label()
      elif label == "comment":
        label_chip.use_comment_label()
      else:
        label_chip.label = label
      search_obj.add_chip(label_chip)

    # Timesketch API returns a max of 10000 results by default
    if search_obj.expected_size > 10000:
      search_obj.max_entries = search_obj.expected_size + 1

    return search_obj.to_pandas()

  def _OutputSearchResults(self, data_frame: pd.DataFrame) -> None:
    """Stores the search results in a container or file.

    Args:
      data_frame: the dataframe containing the Timesketch events.
    """
    if not self.include_internal_columns:
      # Remove internal OpenSearch columns
      data_frame = data_frame.drop(
        columns=["__ts_timeline_id", "_id", "_index", "_source", "_type"],
        errors="ignore",
      )

    if self.output_format == "pandas":
      self.StoreContainer(
        containers.TimesketchEvents(
          name=self.search_name,
          description=self.search_description,
          data_frame=data_frame,
          query=self.query_string,
          sketch_id=self.sketch_id,
        )
      )
    else:
      with tempfile.NamedTemporaryFile(
        mode="w",
        delete=False,
        encoding="utf-8",
        prefix=f"{self.search_name}_" if self.search_name else "",
        suffix=f".{self.output_format}",
      ) as output_file:
        if self.output_format == "csv":
          data_frame.to_csv(output_file, index=False)
        elif self.output_format == "json":
          data_frame.to_json(output_file, orient="records", lines=False)
        elif self.output_format == "jsonl":
          data_frame.to_json(output_file, orient="records", lines=True)
        else:
          self.ModuleError("Unexpected output format", critical=True)
        self.StoreContainer(
          containers.File(
            name=self.search_name,
            description=self.search_description,
            path=output_file.name,
          )
        )

  def FindSketch(self) -> None:
    """Attempts to find a sketch by Sketch ID, cache, or attribute container."""

    if self.sketch:
      return

    if self.sketch_id:
      self.sketch = self._GetSketch(self.sketch_id)
      return

    cached_sketch = self._cache.GetFromCache("timesketch_sketch")
    if cached_sketch:
      self.sketch_id = cached_sketch.id
      self.sketch = cached_sketch
      return

    attributes = self.GetContainers(containers.TicketAttribute)
    self.sketch_id = timesketch_utils.GetSketchIDFromAttributes(attributes)
    if not self.sketch_id:
      return
    self.sketch = self._GetSketch(self.sketch_id)

  def _StoreDataTypesAggregationContainer(
    self, data_frame: pd.DataFrame
  ) -> None:
    """Aggregates the data types in the search results and stores them.

    Args:
      data_frame: the dataframe containing the Timesketch events.
    """
    if "data_type" not in data_frame.columns:
      self.logger.warning(
        "No 'data_type' column found in the search results "
        "skipping aggregation."
      )
      return
    datatypes = data_frame["data_type"].value_counts().to_dict()
    self.StoreContainer(
      containers.TimesketchAggregation(
        name="data_types",
        key="data_type",
        description="Data types in the search results",
        results=datatypes,
      )
    )

  def Process(self) -> None:
    """Processes the Timesketch search query."""
    self.FindSketch()
    if not self.sketch:
      self.ModuleError(
        "Unable to obtain valid sketch ID or sketch, aborting.", critical=True
      )
      return
    data_frame = self._GetSearchResults(
      self.sketch,
      self.query_string,
      self.return_fields,
      self.start_datetime,
      self.end_datetime,
      self.labels,
      self.indices,
    )
    self.logger.info(f"Search returned {len(data_frame)} event(s).")
    if data_frame.empty:
      return
    self._OutputSearchResults(data_frame)
    self._StoreDataTypesAggregationContainer(data_frame)


modules_manager.ModulesManager.RegisterModule(TimesketchSearchEventCollector)
