#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# File:                Ampel-core/ampel/cli/export.py
# License:             BSD-3-Clause
# Author:              valery brinnel <firstname.lastname@gmail.com>
# Date:                16.08.2022
# Last Modified Date:  16.08.2022
# Last Modified By:    valery brinnel <firstname.lastname@gmail.com>

import sys
from collections.abc import Generator
from datetime import datetime
from io import BufferedWriter, TextIOWrapper
from itertools import islice
from json import dumps
from typing import BinaryIO, TextIO

from bson import ObjectId, encode

from ampel.abstract.AbsIdMapper import AbsIdMapper
from ampel.protocol.LoggerProtocol import LoggerProtocol
from ampel.struct.AmpelBuffer import AmpelBuffer
from ampel.util.getch import getch as fgetch
from ampel.util.pretty import prettyjson
from ampel.util.serialize import walk_and_encode

dsi = dict.__setitem__
ufts = datetime.utcfromtimestamp


def txt_export(
	fd: TextIOWrapper | TextIO,
	gen: Generator[AmpelBuffer, None, None],
	id_mapper: None | AbsIdMapper = None,
	chunk_size: int = 200,
	human_times: bool = True,
	pretty: bool = False,
	getch: bool = False,
	close_fd: bool = True,
	logger: None | LoggerProtocol = None
) -> None:

	first = True
	func = prettyjson if pretty else dumps

	fd.write('[\n')

	try:

		while (data := list(islice(gen, chunk_size))):

			walk_and_encode(data)

			if first and len(data) > 0:

				data = iter(data) # type: ignore[assignment]
				el = next(data) # type: ignore[call-overload]
				if id_mapper:
					el['id'] = id_mapper.to_ext_id(el['id'])
				if human_times:
					convert_timestamps(el)
				if logger:
					logger.info(f"Writing content (id: {el['id']})")
				fd.write(func(el))
				first = False

			if getch and fgetch():
				fd.write('\n]\n')
				if logger:
					logger.info("Abording")
				return

			for el in data:

				fd.write(",\n")
				if id_mapper:
					el['id'] = id_mapper.to_ext_id(el['id'])
				if human_times:
					convert_timestamps(el)
				if logger:
					logger.info(f"Writing content (id: {el['id']})")
				fd.write(func(el))

				if getch and fgetch():
					fd.write('\n]\n')
					if logger:
						logger.info("Abording")
					return

		fd.write('\n]\n')

	finally:

		if close_fd:
			fd.flush()
			fd.close()


def bin_export(
	fd: BufferedWriter | BinaryIO,
	gen: Generator[AmpelBuffer, None, None],
	id_mapper: None | AbsIdMapper = None,
	close_fd: bool = True,
	logger: None | LoggerProtocol = None
) -> None:

	for el in gen:
		if logger:
			logger.info(f"Writing content (id: {el['id']})") # type: ignore[str-bytes-safe]
		if id_mapper:
			el['id'] = id_mapper.to_ext_id(el['id'])
		fd.write(encode(el))

	if close_fd:
		if logger:
			logger.info("Closing file descriptor")
		fd.flush()
		fd.close()


def get_fd(
	binary: bool,
	fd: None | str | BufferedWriter | TextIO | TextIOWrapper | BinaryIO,
	close_fd: bool
) -> tuple[BufferedWriter | TextIO | TextIOWrapper | BinaryIO, bool]:

	if binary:
		if isinstance(fd, str):
			return open(fd, "wb"), True
		return fd, close_fd # type: ignore[return-value]
	if fd is None:
		return sys.stdout, False
	if isinstance(fd, str):
		return open(fd, 'w'), close_fd
	return fd, close_fd



def convert_timestamps(ab: AmpelBuffer) -> None:

	try:

		if stock_doc := ab.get('stock'):
			for k in (ts := stock_doc['ts']):
				dsi( # type: ignore[misc]
					ts, k, {
						'tied': ufts(ts[k]['tied']).isoformat(),
						'upd': ufts(ts[k]['upd']).isoformat()
					}
				)
			for je in stock_doc['journal']:
				dsi(je, 'ts', ufts(je['ts']).isoformat()) # type: ignore[arg-type]

		if t1s := ab.get('t1'):
			for t1doc in t1s:
				if isinstance(t1doc['_id'], str) and "oid:" in t1doc['_id']: # type: ignore[typeddict-item]
					gt = ObjectId(t1doc['_id'][4:]).generation_time # type: ignore[typeddict-item]
				else:
					gt = t1doc['_id'].generation_time # type: ignore[typeddict-item]
				dsi(t1doc, 'added', gt.isoformat()) # type: ignore[arg-type]
				for t1meta in t1doc['meta']:
					dsi(t1meta, 'ts', ufts(t1meta['ts']).isoformat()) # type: ignore[arg-type]

		if t2s := ab.get('t2'):
			for t2doc in t2s:
				for t2meta in t2doc['meta']:
					dsi(t2meta, 'ts', ufts(t2meta['ts']).isoformat()) # type: ignore[arg-type]

	except Exception:
		import sys
		import traceback
		traceback.print_exc()
		print(ab)
		sys.exit("Exception occured")
