"""
Module provides tools to deal with Windy's point forecast API.
"""

from dataclasses import dataclass
from datetime import datetime
from enum import Enum


def _json(value):
	try:
		return value.json()
	except AttributeError:
		return value


def _convert_notation(unit):
	return unit.replace("-1", "^-1")


class _StrEnum(Enum):
	def __str__(self):
		return self.value

	def json(self):
		return self.value


class Model(_StrEnum):
	"""
	Numerical models available for use with point forecast API.
	"""
	AROME = "arome"
	GEOS5 = "geos5"
	GFS = "gfs"
	GFSWAVE = "gfsWave"
	ICONEU = "iconEu"
	NAMALASKA = "namAlaska"
	NAMCONUS = "namConus"
	NAMHAWAII = "namHawaii"


class Level(_StrEnum):
	"""
	Selectable levels for some of the input parameters that support them.
	"""
	SURFACE = "surface"
	H1000 = "1000h"
	H950 = "950h"
	H925 = "925h"
	H900 = "900h"
	H850 = "850h"
	H800 = "800h"
	H700 = "700h"
	H600 = "600h"
	H500 = "500h"
	H400 = "400h"
	H300 = "300h"
	H200 = "200h"
	H150 = "150h"


@dataclass
class Request:
	"""
	Wraps raw JSON request expected by Windy's API.
	"""
	key: str
	lat: float
	lon: float
	model: Model
	parameters: list = None
	levels: list = None

	def json(self):
		body = {
			'key': self.key,
			'lat': self.lat,
			'lon': self.lon,
			'model': _json(self.model),
			'parameters': self.parameters or [],
		}
		if self.levels:
			body['levels'] = [_json(x) for x in self.levels]
		return body


class EntryView:
	"""
	Allows to iterate over samples in Response in a zip-like manner, where all parameters for a given time point are
	available via item access.
	"""
	def __init__(self, response, index=0):
		self._response = response
		self._index = index

	def __iter__(self):
		return self

	def __next__(self):
		self._index += 1
		if not self:
			raise StopIteration
		return self

	def __bool__(self):
		return self._index < len(self)

	def __len__(self):
		return len(self._response)

	@property
	def timestamp(self):
		"""
		Datetime object representing timestamp of the current entry.
		"""
		return self._response.timestamps[self._index]

	def __getitem__(self, key):
		return self._response.samples[key][self._index]


class Response:
	"""
	Wraps raw JSON response from the Windy's API to allow for easier access, converts all values to pint's
	Quantities, and converts all timestamps into datetime objects.

	Can be used in a for-loop to access all samples via EntryView:

		>>> for entry in response:
		>>>     print(entry.timestamp, entry['temp-surface'])

	Otherwise, timestamps list and samples dictionary are available for direct access.
	"""
	_INTERNAL_FIELDS = ('ts', 'units', 'warning')

	def __init__(self, registry, raw):
		self.timestamps = [datetime.fromtimestamp(x // 1000) for x in raw['ts']]
		self.samples = {}
		parameters = ((x, raw['units'][x]) for x in raw if x not in self._INTERNAL_FIELDS)
		for parameter, unit in parameters:
			self.samples[parameter] = [x * registry(_convert_notation(unit)) for x in raw[parameter]]

	def __len__(self):
		return len(self.timestamps)

	def parameters(self) -> tuple:
		"""
		All of the available output parameters.
		"""
		return tuple(self.samples.keys())

	def entries(self) -> EntryView:
		"""
		Helper iterator to go over all of the samples in a zip-like manner.
		"""
		return EntryView(self)

	def __iter__(self):
		return self.entries()


@dataclass
class PointForecast:
	"""
	Represents the point forecast endpoint bound to *path*. Once created it can be called with Request object or
	with the same arguments that would be used to initialize the Request. The request is made using the passed
	*ctx*, which is usually a Windy instance.
	"""
	path: str

	def __call__(self, ctx, *args, **kwargs):
		try:
			body = args[0].json()
		except (IndexError, AttributeError):
			body = Request(*args, **kwargs).json()
		response = ctx.session.post(ctx.api + self.path, json=body)
		response.raise_for_status()
		return Response(ctx.registry, response.json())