# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Builds Creative Maps network."""

# pylint: disable=C0330, g-bad-import-order, g-multiple-import

from __future__ import annotations

import base64
import json
import logging
import operator
import os
import textwrap
from collections.abc import Mapping, Sequence
from io import BytesIO
from typing import Any, TypeAlias, TypedDict

import numpy as np
import pydantic
import smart_open
from garf_core import report
from media_similarity import media_similarity_service
from media_tagging import media, tagging_result
from PIL import Image, ImageDraw, ImageFont, ImageOps

MetricInfo: TypeAlias = dict[str, int | float]
Info: TypeAlias = dict[str, int | float | str | list[str] | None]


class MediaInfo(pydantic.BaseModel):
  """Contains extra information on a given medium."""

  media_path: str
  media_name: str
  info: Info
  series: dict[str, MetricInfo]
  media_preview: str | None = None
  size: float | None = None
  segments: dict[str, dict[str, Info]] | None = None

  def model_post_init(self, __context__) -> None:  # noqa: D105
    if not self.media_preview:
      self.media_preview = self.media_path
    self.info = dict(self.info)


class GraphInfo(TypedDict):
  adaptive_threshold: float
  period: dict[str, str]


class ClusterInfo(TypedDict):
  name: str


class NodeInfo(TypedDict):
  name: str
  label: str
  type: str
  image: str
  media_path: str
  cluster: int
  info: Info
  series: dict[str, MetricInfo]
  tags: list[dict[str, float]]
  segments: dict[str, str]


class CreativeMapJson(TypedDict):
  graph: GraphInfo
  clusters: dict[int, ClusterInfo]
  nodes: list[NodeInfo]
  edges: list[dict[str, int | float]]


class CreativeMap:
  """Defines CreativeMap based on a graph.

  Attributes:
    adaptive_threshold: Minimal value for defining similar media.
    fetching_request: Additional parameter used to generate a map.
    nodes: Information on each node of the map.
    edges: Information on each edge of the map.
    clusters: Aggregated information on each cluster of the map.
  """

  def __init__(
    self,
    adaptive_threshold: float,
    fetching_request: dict[str, Any] | None = None,
    nodes: list[NodeInfo] | None = None,
    edges: list[dict[str, int | float]] | None = None,
    clusters: dict[int, ClusterInfo] | None = None,
  ) -> None:
    """Initializes CreativeMap."""
    self.adaptive_threshold = adaptive_threshold
    self.fetching_request = fetching_request or {}
    self.nodes: list[NodeInfo] = nodes or []
    self.edges: list[dict[str, int | float]] = edges or []
    self.clusters: dict[int, ClusterInfo] = clusters or {}

  @classmethod
  def from_clustering(
    cls,
    clustering_results: media_similarity_service.ClusteringResults,
    tagging_results: Sequence[tagging_result.TaggingResult],
    extra_info: dict[str, MediaInfo] | None = None,
    fetching_request: dict[str, Any] | None = None,
    embed_previews: bool = False,
  ) -> CreativeMap:
    """Builds network visualization with injected extra_info."""
    tagging_mapping = {
      result.identifier: result.content for result in tagging_results
    }
    tagging_hash_to_identifier_mapping = {
      t.hash: t.identifier for t in tagging_results
    }
    media_type = tagging_results[0].type
    for node in clustering_results.graph.nodes:
      node_hash = node.get('name', '')
      node_name = tagging_hash_to_identifier_mapping.get(node_hash)
      if node_extra_info := extra_info.get(node_name):
        node['id'] = node_hash
        if size := node_extra_info.size:
          node['size'] = size
        node['type'] = media_type
        if media_type == 'text':
          node['image'] = _create_text_image_bytes(node_extra_info.media_path)
        elif embed_previews:
          node['image'] = _embed_preview(node_extra_info.media_preview)
        else:
          node['image'] = node_extra_info.media_preview
        node['media_path'] = node_extra_info.media_path
        node['label'] = node_extra_info.media_name or 'Unknown'
        node['cluster'] = clustering_results.clusters.get(node_name)
        node['info'] = node_extra_info.info
        node['series'] = node_extra_info.series
        node['tags'] = [
          {'tag': tag.name.replace("'", ''), 'score': tag.score}
          for tag in tagging_mapping.get(node_name, [])
        ]
        node['segments'] = node_extra_info.segments
    clusters = {
      cluster_id: f'Cluster: {cluster_id}'
      for cluster_id in set(clustering_results.clusters.values())
    }
    edges = [
      {
        'from': _from,
        'to': to,
        'similarity': similarity,
      }
      for _from, to, similarity in clustering_results.graph.edges
    ]
    return CreativeMap(
      adaptive_threshold=clustering_results.adaptive_threshold,
      fetching_request=fetching_request,
      nodes=clustering_results.graph.nodes,
      edges=edges,
      clusters=clusters,
    )

  def to_json(self) -> CreativeMapJson:
    """Extracts nodes from Network."""
    start_date = ''
    end_date = ''
    for node in self.nodes:
      if series := node.get('series'):
        min_start_date = min(series.keys())
        max_end_date = max(series.keys())
        if not start_date or min_start_date < start_date:
          start_date = min_start_date
        if not end_date or max_end_date > end_date:
          end_date = max_end_date

    return {
      'graph': {
        'adaptive_threshold': self.adaptive_threshold,
        'period': {
          'start_date': start_date or 'null',
          'end_date': end_date or 'null',
        },
      },
      'clusters': self.clusters,
      'nodes': self.nodes,
      'edges': self.edges,
    }

  def save(self, path: str | os.PathLike[str]) -> None:
    """Saves map to a json file."""
    with smart_open.open(path, 'w', encoding='utf-8') as f:
      json.dump(self.to_json(), f)


def convert_report_to_media_info(
  performance: report.GarfReport,
  media_type: media.MediaTypeEnum,
  metric_columns: Sequence[str] | None = None,
  segment_columns: Sequence[str] | None = None,
  modules: Sequence[str] | None = None,
  with_size_base: str | None = None,
) -> dict[str, MediaInfo]:
  """Convert report to MediaInfo mappings.

  Args:
    performance: Report with media data.
    media_type: Type of media in the report.
    metric_columns:  Name of metrics to be included in media info.
    segment_columns:  Name of segments to be calculated.
    modules: Optional column names to be added to media_info.
    with_size_base: Optional column name for regulating size of media info.

  Returns:
    Mapping between media identifier and its media info.
  """
  if isinstance(media_type, str):
    media_type = media.MediaTypeEnum[media_type.upper()]
  if with_size_base and with_size_base not in performance.column_names:
    logging.warning('Failed to set MediaInfo size to %s', with_size_base)
    with_size_base = None
  if with_size_base:
    try:
      float(performance[0][with_size_base])
    except TypeError:
      logging.warning('MediaInfo size attribute should be numeric')
    with_size_base = None

  if modules:
    modules = {module.split('.')[1] for module in modules}
    modules = modules.intersection(performance.column_names)
  else:
    modules = set()

  if media_type == media.MediaTypeEnum.YOUTUBE_VIDEO:
    modules.add('duration')
    modules.add('orientation')
  elif media_type == media.MediaTypeEnum.IMAGE:
    modules.add('format_type')

  metric_columns = metric_columns or []
  if 'in_campaigns' in performance.column_names:
    modules.add('in_campaigns')

  results = {}
  for media_url, values in performance.to_dict(key_column='media_url').items():
    info = build_info(values, list(metric_columns), list(modules))
    segments = (
      build_segments(values, segment_columns, metric_columns)
      if segment_columns
      else {}
    )
    if values[0].get('date'):
      series = build_segments(values, ['date'], metric_columns).get('date', {})
    else:
      series = {}
    if with_size_base and (size_base := info.get(with_size_base)):
      media_size = np.log(size_base) * np.log10(size_base)
    else:
      media_size = None
    results[media.convert_path_to_media_name(media_url, media_type)] = (
      MediaInfo(
        **create_node_links(media_url, media_type),
        media_name=values[0].get('media_name'),
        info=info,
        series=series,
        size=media_size,
        segments=segments,
      )
    )
  return results


def build_info(
  data: Info, metric_names: Sequence[str], modules: Sequence[str] | None = None
) -> Info:
  """Extracts and aggregated data for specified metrics.

  Args:
    data: All available information on media_url.
    metric_names: Metrics to sum.
    modules: Meterics/ dimensions to get as is without aggregation.

  Returns:
    Mapping between name of metric / dimension and its processed value.
  """
  info = {
    metric: _aggregate_nested_metric(data, metric) for metric in metric_names
  }
  if modules:
    dimensions = {module: data[0].get(module) for module in modules}
    info.update(dimensions)
  return info


def build_segments(
  data: Info, segment_names: Sequence[str], metric_names: Sequence[str]
) -> dict[str, dict[Info]]:
  """Builds info object for each variant of a segment.

  Args:
    data: Report data formatted as a mapping.
    segment_names: Names of column in report to transform into segments.
    metric_names: Names of metrics in report to calculate for each segment.

  Returns:
    Mapping between each segment, it's variants and corresponding metrics.
  """
  segments = {}
  for segment_name in segment_names:
    get_segment_getter = operator.itemgetter(segment_name)
    try:
      segment_values = set(map(get_segment_getter, data))
    except KeyError:
      continue
    segment_variants = {}
    for segment_value in segment_values:
      if segment_value != 'UNKNOWN':
        segment_variants[segment_value] = build_info(
          list(filter(lambda x: x[segment_name] == segment_value, data)),
          metric_names,
        )
    if segment_variants:
      segments[segment_name] = segment_variants
  return segments


def _aggregate_nested_metric(
  data: Info | Sequence[Info],
  metric_name: str,
) -> float | int | str | list[str] | None:
  """Performance appropriate aggregation over a dictionary.

  Sums numerical values and deduplicates and sorts alphabetically
  string values.

  Args:
    data: Data to extract metrics from.
    metric_name: Name of a metric to be extracted from supplied data.

  Returns:
    Aggregated value of a metric.
  """
  get_metric_getter = operator.itemgetter(metric_name)
  if isinstance(data, Mapping):
    return get_metric_getter(data)

  try:
    res = list(map(get_metric_getter, data))
  except KeyError:
    return None
  try:
    return sum(res)
  except TypeError:
    if len(result := sorted(set(res))) == 1:
      return ','.join(result)
    return result


def create_node_links(
  url: str, media_type: media.MediaTypeEnum
) -> dict[str, str]:
  return {
    'media_path': _to_youtube_video_link(url)
    if media_type == media.MediaTypeEnum.YOUTUBE_VIDEO
    else url,
    'media_preview': _to_youtube_preview_link(url)
    if media_type == media.MediaTypeEnum.YOUTUBE_VIDEO
    else url,
  }


def _to_youtube_preview_link(video_id: str) -> str:
  return f'https://img.youtube.com/vi/{video_id}/0.jpg'


def _to_youtube_video_link(video_id: str) -> str:
  return f'https://www.youtube.com/watch?v={video_id}'


def _create_text_image_bytes(
  text: str,
  format: str = 'PNG',
  font_size: int = 40,
  font_path: str | None = None,
  trim_treshold: int = 100,
  border=1,
):
  text = text[:trim_treshold]
  text = '\n'.join(textwrap.wrap(text, width=20))
  dummy_img = Image.new(mode='RGB', size=(1, 1))
  dummy_draw = ImageDraw.Draw(dummy_img)

  font = (
    ImageFont.truetype(font_path, font_size)
    if font_path
    else ImageFont.load_default(size=font_size)
  )

  bbox = dummy_draw.textbbox((0, 0), text, font=font)
  width = bbox[2] - bbox[0] + 10
  height = bbox[3] - bbox[1] + 10

  img = Image.new(mode='RGB', size=(width, height), color='white')
  img = ImageOps.expand(img, border=border, fill=(211, 211, 211))
  draw = ImageDraw.Draw(img)

  draw.multiline_text((5, 5), text, fill='black', font=font)

  byte_stream = BytesIO()
  img.save(byte_stream, format=format)

  image_bytes = byte_stream.getvalue()
  encoded_image = base64.b64encode(image_bytes).decode('utf-8')
  return f'data:image/png;base64,{encoded_image}'


def _embed_preview(url: str) -> str:
  logging.info('Embedding preview for url %s', url)
  with smart_open.open(url, 'rb') as f:
    image_preview = f.read()
  encoded_image = base64.b64encode(image_preview).decode('utf-8')
  return f'data:image/png;base64,{encoded_image}'
