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

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

import numpy as np


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


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


class _StrEnum(Enum):
	def __init__(self, value):
		self._index = len(self.__class__.__members__)

	def __lt__(self, other):
		if self.__class__ is other.__class__:
			return self._index < other._index
		return NotImplemented

	def __le__(self, other):
		if self.__class__ is other.__class__:
			return self._index <= other._index
		return NotImplemented

	def __gt__(self, other):
		if self.__class__ is other.__class__:
			return self._index > other._index
		return NotImplemented

	def __ge__(self, other):
		if self.__class__ is other.__class__:
			return self._index >= other._index
		return NotImplemented

	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"

	def pressure(self):
		if self is Level.SURFACE:
			raise ValueError
		return float(self.value[:-1])


@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 Prediction:
	"""
	Predicted values for each of the requested parameters along with their associated time point.
	"""
	def __init__(self, response, index=0):
		self._response = response
		self._index = index

	@property
	def timestamp(self) -> datetime:
		return self._response.timestamps[self._index]

	@property
	def parameters(self) -> tuple:
		return self._response.parameters

	@property
	def levels(self) -> tuple:
		return self._response.levels

	def __iter__(self):
		return iter(self.parameters)

	def __getitem__(self, key):
		return self._response.values[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 Prediction:

		>>> for prediction in response:
		>>>     print(prediction.timestamp, prediction['temp'])

	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']]
		count = len(self.timestamps)
		split = [tuple(x.split("-")) for x in raw if x not in self._INTERNAL_FIELDS]
		self.parameters = tuple({x for x, _ in split})
		self.levels = tuple(sorted({Level(x) for _, x in split}))
		if 'pressure' in self.parameters:
			for level in self.levels:
				if level is Level.SURFACE:
					continue
				raw[f'pressure-{level}'] = [level.pressure() for _ in range(len(self.timestamps))]
		units = {x: registry(_convert_notation(raw['units'][f'{x}-surface'])) for x in self.parameters}  # Don't guess surface
		self.values = {p: [[raw[f'{p}-{l}'][i] for l in self.levels] * units[p] for i in range(count)] for p in self.parameters}

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

	def predictions(self) -> Prediction:
		"""
		Yields Prediction for each time point available in this Response.
		"""
		for index in range(len(self.timestamps)):
			yield Prediction(self, index)

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


@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())