#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File:                Ampel-core/ampel/t3/supply/T3DefaultBufferSupplier.py
# License:             BSD-3-Clause
# Author:              valery brinnel <firstname.lastname@gmail.com>
# Date:                14.07.2021
# Last Modified Date:  13.12.2021
# Last Modified By:    valery brinnel <firstname.lastname@gmail.com>

from collections.abc import Generator, Sequence

from ampel.abstract.AbsBufferComplement import AbsBufferComplement
from ampel.abstract.AbsT3Loader import AbsT3Loader
from ampel.abstract.AbsT3Selector import AbsT3Selector
from ampel.abstract.AbsT3Supplier import AbsT3Supplier
from ampel.model.UnitModel import UnitModel
from ampel.struct.AmpelBuffer import AmpelBuffer
from ampel.struct.T3Store import T3Store
from ampel.util.collections import get_chunks as chunks_func


class T3DefaultBufferSupplier(AbsT3Supplier[Generator[AmpelBuffer, None, None]]):
	"""
	Default T3 supplier that:
	- matches stock ids based on the stock collection (section 'select'),
	- loads AmpelBuffer content as defined in the 'load' section
	- potentially updates AmpelBuffer through the 'complement' section
	"""

	#: Select stocks
	select: UnitModel

	#: Fill :class:`~ampel.core.AmpelBuffer.AmpelBuffer` for each selected stock
	load: UnitModel

	#: Add external information to each :class:`~ampel.core.AmpelBuffer.AmpelBuffer`.
	complement: None | Sequence[UnitModel]

	#: number of stocks to load at once. Set to 0 to disable chunking
	chunk_size: int = 1000


	def __init__(self, **kwargs) -> None:

		super().__init__(**kwargs)

		# stock selection
		#################

		# Spawn and run a new selector instance
		# fetch() returns an iterable (often a pymongo cursor)
		self.selector = self.context.loader.new_context_unit(
			model = self.select,
			context = self.context,
			sub_type = AbsT3Selector,
			logger = self.logger
		)

		# Content loader
		################

		# Spawn requested content loader
		self.data_loader = self.context.loader \
			.new_context_unit(
				model = self.load,
				context = self.context,
				sub_type = AbsT3Loader,
				logger = self.logger
			)

		# Content complementer
		######################

		if self.complement:

			# Spawn requested snapdata complementers
			self.complementers: None | list[AbsBufferComplement] = [
				self.context.loader \
					.new_context_unit(
						model = conf_el,
						context = self.context,
						sub_type = AbsBufferComplement,
						logger = self.logger
					)
				for conf_el in self.complement
			]
		else:
			self.complementers = None


	def supply(self, t3s: T3Store) -> Generator[AmpelBuffer, None, None]:

		# NB: we consume the entire cursor at once using list() to be robust
		# against cursor timeouts or server restarts during long lived T3 processes
		stock_ids = list(self.selector.fetch() or [])
		if not stock_ids:
			return

		# Usually, id_key is '_id' but it can be 'stock' if the
		# selection is based on t2 documents for example
		id_key = self.selector.field_name

		# Run start
		###########
		chunks = chunks_func(stock_ids, self.chunk_size) if self.chunk_size > 0 else [stock_ids]

		# Loop over chunks from the cursor/iterator
		for chunk_ids in chunks:

			# allow working chunks to complete even if some raise exception
			try:

				# Load info from DB
				tran_data = self.data_loader.load([sid[id_key] for sid in chunk_ids])

				# Potentialy add complementary information (spectra, TNS names, ...)
				if self.complementers:
					for appender in self.complementers:
						appender.complement(tran_data, t3s)

				yield from tran_data

			except Exception as e:  # noqa: PERF203
				self.event_hdlr.handle_error(e, self.logger)
