# coding=utf-8
# Copyright 2026 The Perch Authors.
#
# 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
#
#     http://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.

"""Functionality for embedding audio examples."""

from concurrent import futures
import dataclasses
import functools
import itertools
import threading

from absl import logging
import audioread
from ml_collections import config_dict
import numpy as np
from perch_hoplite import audio_io
from perch_hoplite.agile import metadata
from perch_hoplite.agile import source_info
from perch_hoplite.db import datatypes
from perch_hoplite.db import interface as hoplite_interface
from perch_hoplite.zoo import model_configs
from perch_hoplite.zoo import zoo_interface
import soundfile


@dataclasses.dataclass
class ModelConfig(datatypes.HopliteConfig):
  """Configuration for embedding model.

  Attributes:
    model_key: Key for the model wrapper class.
    embedding_dim: Dimensionality of the embedding.
    model_config: Config dict of arguments to instantiate the model wrapper.
    logit_key: If provided, model predictions will be stored instead of raw
      embeddings.
    logit_idxes: When storing model predictions, allows selecting a subset of
      prediction classes.
  """

  model_key: str
  embedding_dim: int
  model_config: config_dict.ConfigDict
  logits_key: str | None = None
  logits_idxes: tuple[int, ...] | None = None


def worker_initializer(state):
  name = threading.current_thread().name
  state[name + 'db'] = state['db'].thread_split()


def process_source_id(
    state,
    source_id: source_info.SourceId,
    window_size_s: float,
):
  """Process a single audio source."""
  worker = state['worker']
  # To access the thread-specific DB, we need to use the thread name.
  # name = threading.current_thread().name
  # db = state[name + 'db']
  glob = worker.audio_globs[source_id.dataset_name]
  target_sample_rate = worker.get_sample_rate_hz(source_id)
  audio_array = worker.load_audio(source_id)

  if audio_array is None:
    return
  if (
      audio_array.shape[0]
      < glob.min_audio_len_s * worker.embedding_model.sample_rate
  ):
    return

  outputs = worker.embedding_model.embed(audio_array)
  logits_key = state['worker'].model_config.logits_key
  if logits_key is None:
    embeddings = outputs.embeddings
  else:
    embeddings = outputs.logits[logits_key]
    logits_idxes = state['worker'].model_config.logits_idxes
    if logits_idxes is not None:
      embeddings = embeddings[..., logits_idxes]
    # Add channel axis to match the expected shape of the embeddings.
    embeddings = embeddings[:, np.newaxis, :]

  if embeddings is None:
    return

  sources = []
  offsets = []
  embs = []

  hop_size_s = worker.compute_hop_size_s(source_id, target_sample_rate)
  for t, embedding in enumerate(embeddings):
    offset_s = source_id.offset_s + t * hop_size_s
    offsets_list = [offset_s, offset_s + window_size_s]
    for channel_embedding in embedding:
      sources.append(source_id)
      offsets.append(offsets_list)
      embs.append(channel_embedding)
  return sources, offsets, embs


# TODO(tomdenton): Use itertools.batched in Python 3.12+
def batched(iterable, n):
  it = iter(iterable)
  while batch := tuple(itertools.islice(it, n)):
    yield batch


class EmbedWorker:
  """Worker for embedding audio examples."""

  def __init__(
      self,
      audio_sources: source_info.AudioSources,
      model_config: ModelConfig,
      db: hoplite_interface.HopliteDBInterface,
      embedding_model: zoo_interface.EmbeddingModel | None = None,
      audio_worker_threads: int = 8,
  ):
    self.db = db
    self.model_config = model_config
    self.audio_sources = audio_sources
    self.audio_worker_threads = audio_worker_threads
    if embedding_model is None:
      model_class = model_configs.get_model_class(model_config.model_key)
      self.embedding_model = model_class.from_config(model_config.model_config)
    else:
      self.embedding_model = embedding_model
    self.window_size_s = getattr(self.embedding_model, 'window_size_s')
    self.audio_globs = {
        g.dataset_name: g for g in self.audio_sources.audio_globs
    }
    self.metadata = {
        g.dataset_name: metadata.AgileMetadata.from_directory(g.base_path)
        for g in self.audio_sources.audio_globs
    }

  def _log_error(self, source_id, exception, counter_name):
    logging.warning(
        'The audio at (%s / %f) could not be loaded (%s). '
        'The exception was (%s)',
        source_id.filepath,
        source_id.offset_s,
        counter_name,
        exception,
    )

  def _update_audio_sources(self):
    """Validates the embed config and/or saves it to the DB."""
    db_metadata = self.db.get_metadata(None)
    if 'audio_sources' not in db_metadata:
      self.db.insert_metadata(
          'audio_sources', self.audio_sources.to_config_dict()
      )
      return

    db_audio_sources = source_info.AudioSources.from_config_dict(
        db_metadata['audio_sources']
    )
    merged = self.audio_sources.merge_update(db_audio_sources)
    self.db.insert_metadata('audio_sources', merged.to_config_dict())
    self.audio_sources = merged

  def _update_model_config(self):
    """Validates the model config and/or saves it to the DB."""
    db_metadata = self.db.get_metadata(None)
    if 'model_config' not in db_metadata:
      self.db.insert_metadata(
          'model_config', self.model_config.to_config_dict()
      )
      return

    db_model_config = ModelConfig(**db_metadata['model_config'])
    if self.model_config == db_model_config:
      return

    # Validate the config against the DB.
    # TODO(tomdenton): Implement compatibility checks for model configs.
    if self.model_config.model_key != db_model_config.model_key:
      raise AssertionError(
          'The configured model key does not match the model key that is '
          'already in the DB.'
      )
    if self.model_config.embedding_dim != db_model_config.embedding_dim:
      raise AssertionError(
          'The configured embedding dimension does not match the embedding '
          'dimension that is already in the DB.'
      )
    self.db.insert_metadata('model_config', self.model_config.to_config_dict())

  def update_configs(self):
    """Validates the configs and saves them to the DB."""
    self._update_model_config()
    self._update_audio_sources()
    self.db.commit()

  @functools.cache  # pylint: disable=method-cache-max-size-none
  def _get_or_insert_deployment_id(
      self,
      deployment_name: str,
      project_name: str,
      error_on_insert: bool = False,
  ) -> int:
    """Get the deployment ID for the given deployment name and project name."""
    deployments = self.db.get_all_deployments(
        config_dict.create(eq=dict(name=deployment_name, project=project_name))
    )
    if deployments:
      return deployments[0].id
    if error_on_insert:
      raise ValueError(
          f'Deployment {deployment_name} not found in project {project_name}.'
      )
    md = self.metadata[project_name].get_deployment_metadata(deployment_name)
    md.pop('deployment', None)
    return self.db.insert_deployment(
        name=deployment_name,
        project=project_name,
        **md,
    )

  @functools.cache  # pylint: disable=method-cache-max-size-none
  def _get_or_insert_recording_id(
      self,
      filename: str,
      deployment_id: int,
      dataset_name: str,
      error_on_insert: bool = False,
  ) -> tuple[int, bool]:
    """Get the recording ID, and indicate whether it was newly inserted."""
    recordings = self.db.get_all_recordings(
        config_dict.create(
            eq=dict(filename=filename, deployment_id=deployment_id)
        )
    )
    if recordings:
      return recordings[0].id, False
    if error_on_insert:
      raise ValueError(
          f'Recording {filename} not found in deployment {deployment_id}.'
      )
    md = self.metadata[dataset_name].get_recording_metadata(filename)
    md.pop('recording', None)
    return (
        self.db.insert_recording(
            filename=filename,
            deployment_id=deployment_id,
            **md,
        ),
        True,
    )

  def add_deployments(self, target_dataset_name: str | None = None):
    """Add deployments to db and create a source ID to deployment ID mapping."""
    # Create missing deployments in the database.
    for source in self.audio_sources.iterate_all_sources(target_dataset_name):
      deployment_name = source.deployment_name_from_file_id()
      self._get_or_insert_deployment_id(
          deployment_name=deployment_name,
          project_name=source.dataset_name,
      )
    self.db.commit()

  def add_recordings(self, target_dataset_name: str | None = None) -> set[int]:
    """Add recordings to db and create a source ID to recording ID mapping."""
    new_recordings = set([])
    for source in self.audio_sources.iterate_all_sources(target_dataset_name):
      deployment_id = self._get_or_insert_deployment_id(
          deployment_name=source.deployment_name_from_file_id(),
          project_name=source.dataset_name,
      )
      recording_id, is_new = self._get_or_insert_recording_id(
          source.file_id,
          deployment_id,
          source.dataset_name,
      )
      if is_new:
        new_recordings.add(recording_id)
    self.db.commit()
    return new_recordings

  def add_annotations(self, target_dataset_name: str | None = None):
    """Add annotations from metadata to db."""
    dataset_names = self.metadata.keys()
    if target_dataset_name is not None:
      dataset_names = [target_dataset_name]

    for dataset_name in dataset_names:
      if dataset_name not in self.metadata:
        continue
      agile_md = self.metadata[dataset_name]
      if not agile_md.annotations:
        continue
      for file_id, annotation_list in agile_md.annotations.items():
        depl_name = file_id.split('/')[0]
        if not depl_name:
          logging.warning(
              'Could not get deployment name from file_id %s, skipping.',
              file_id,
          )
          continue
        depl_id = self._get_or_insert_deployment_id(depl_name, dataset_name)
        rec_id, _ = self._get_or_insert_recording_id(
            file_id, depl_id, dataset_name
        )
        for annotation in annotation_list:
          self.db.insert_annotation(
              rec_id,
              annotation.offsets,
              annotation.label,
              annotation.label_type,
              provenance=annotation.provenance,
              handle_duplicates='skip',
          )
    self.db.commit()

  def embed_dataset(
      self,
      batch_size=32,
      handle_duplicates: str = 'error',
      target_dataset_name: str | None = None,
      new_recordings: set[int] | None = None,
  ):
    """Embed audio examples from the given dataset."""
    # Process all sources.
    state = {}
    state['db'] = self.db
    state['worker'] = self
    state['new_recordings'] = new_recordings
    with futures.ThreadPoolExecutor(
        max_workers=self.audio_worker_threads,
        initializer=worker_initializer,
        initargs=(state,),
    ) as executor:
      source_iterator = self.audio_sources.iterate_all_sources(
          target_dataset_name
      )
      for source_ids_batch in batched(source_iterator, batch_size):
        got = executor.map(
            process_source_id,
            itertools.repeat(state),
            source_ids_batch,
            itertools.repeat(self.window_size_s),
        )
        # TODO(tomdenton): Consider using a db writer thread to avoid blocking.
        for result in got:
          if result is None:
            continue
          recording_ids = []
          for s in result[0]:
            deployment_id = self._get_or_insert_deployment_id(
                s.deployment_name_from_file_id(), s.dataset_name
            )
            recording_id, _ = self._get_or_insert_recording_id(
                s.file_id, deployment_id, s.dataset_name
            )
            recording_ids.append(recording_id)
          if all(r in new_recordings for r in recording_ids):
            dupe_strategy = 'allow'
          else:
            dupe_strategy = handle_duplicates
          windows_batch = [
              {
                  'recording_id': recording_ids[i],
                  'offsets': o,
              }
              for i, (_, o, _) in enumerate(zip(*result))
          ]
          embeddings_batch = np.array(result[2])
          self.db.insert_windows_batch(
              windows_batch,
              embeddings_batch,
              handle_duplicates=dupe_strategy,
          )
    self.db.commit()

  def get_sample_rate_hz(self, source_id: source_info.SourceId) -> int:
    """Get the sample rate of the embedding model."""
    dataset_name = source_id.dataset_name
    if dataset_name not in self.audio_globs:
      raise ValueError(f'Dataset name {dataset_name} not found in audio globs.')
    audio_glob = self.audio_globs[dataset_name]
    if audio_glob.target_sample_rate_hz == -2:
      return self.embedding_model.sample_rate
    elif audio_glob.target_sample_rate_hz == -1:
      # Uses the file's native sample rate.
      return -1
    elif audio_glob.target_sample_rate_hz > 0:
      return audio_glob.target_sample_rate_hz
    else:
      raise ValueError('Invalid target_sample_rate.')

  def load_audio(self, source_id: source_info.SourceId) -> np.ndarray | None:
    """Load audio from the indicated source and log any problems."""
    target_sample_rate_hz = self.get_sample_rate_hz(source_id)
    try:
      audio_array = audio_io.load_audio_window(
          filepath=source_id.filepath,
          offset_s=source_id.offset_s,
          sample_rate=target_sample_rate_hz,
          window_size_s=source_id.shard_len_s,
      )
      return np.array(audio_array)
    except soundfile.LibsndfileError as inst:
      self._log_error(source_id, inst, 'audio_libsndfile_error')
    except ValueError as inst:
      self._log_error(source_id, inst, 'audio_bad_offset')
    except audioread.NoBackendError as inst:
      self._log_error(source_id, inst, 'audio_no_backend')
    except EOFError as inst:
      self._log_error(source_id, inst, 'audio_eof_error')
    except RuntimeError as inst:
      if 'Soundfile is not available' in str(inst):
        self._log_error(source_id, inst, 'audio_no_soundfile')
      else:
        self._log_error(source_id, inst, 'audio_runtime_error')

  def compute_hop_size_s(
      self,
      source_id: source_info.SourceId,
      target_sample_rate_hz: int,
      model_hop_size_s: float | None = None,
  ) -> float:
    """Compute the hop size of the embedding model."""

    if model_hop_size_s is not None:
      return model_hop_size_s
    if hasattr(self.embedding_model, 'hop_size_s'):
      model_hop_size_s = float(self.embedding_model.hop_size_s)
    else:
      # TODO(tomdenton): Allow user specified hop size.
      raise ValueError('hop_size_s is not defined for the model.')
    model_sample_rate = self.embedding_model.sample_rate

    if target_sample_rate_hz == -2:
      return model_hop_size_s
    elif target_sample_rate_hz == -1:
      audio_sample_rate = source_id.sample_rate_hz
    elif target_sample_rate_hz > 0:
      audio_sample_rate = target_sample_rate_hz
    else:
      raise ValueError('Invalid target_sample_rate.')
    return model_hop_size_s * model_sample_rate / audio_sample_rate

  def process_all(
      self,
      target_dataset_name: str | None = None,
      batch_size=32,
      handle_duplicates='error',
  ):
    """Process all audio examples."""

    # Update model config and audio sources in the database.
    self.update_configs()
    if self.db.count_embeddings() == 0:
      # No chance of duplicates, so we can use "allow" mode.
      handle_duplicates = 'allow'

    # Add deployments and recordings to the database.
    print('\nAdding deployments...')
    self.add_deployments(target_dataset_name)
    print('\nAdding recordings...')
    new_recordings = self.add_recordings(target_dataset_name)
    print('\nAdding annotations...')
    self.add_annotations()
    print('\nEmbedding audio...')
    self.embed_dataset(
        batch_size=batch_size,
        handle_duplicates=handle_duplicates,
        target_dataset_name=target_dataset_name,
        new_recordings=new_recordings,
    )
    self.db.commit()
