"""Functions for writing to pvsite db."""

import datetime as dt
import logging
import uuid

import dateutil.parser as dp
import numpy.typing as npt
import pandas as pd
import sqlalchemy.orm as sa_orm

from pvsite_datamodel.read.site import get_site_by_uuid
from pvsite_datamodel.sqlmodels import ForecastSQL, ForecastValueSQL
from pvsite_datamodel.write.datetime_intervals import get_or_else_create_datetime_interval
from pvsite_datamodel.write.upsert import upsert
from pvsite_datamodel.write.utils import WrittenRow


def insert_forecast_values(
    session: sa_orm.Session, df_forecast_values: pd.DataFrame, forecast_version: str = "0.0.0"
) -> list[WrittenRow]:
    """Insert a dataframe of forecast values into the database.

    :param session: sqlalchemy session for interacting with the database
    :param df_forecast_values: pandas dataframe with columns
    ["target_datetime_utc", "forecast_kw", "pv_uuid"]
    :param forecast_version: the version of the model used to create the forecast
    :return list[WrittenRow]: list of added rows to DB
    """
    # Track rows added to DB
    written_rows: list[WrittenRow] = []

    # Loop over all the unique sites that have got forecast values
    sites: npt.ndarray[uuid.UUID] = df_forecast_values["pv_uuid"].unique()
    sql_forecast_values = []
    for site_uuid in sites:

        # Check whether the site id exits in the table, otherwise return an error
        get_site_by_uuid(session=session, site_uuid=site_uuid)

        # Create a forcast (sequence) for the site, and write it to db
        forecast: ForecastSQL = ForecastSQL(
            forecast_uuid=uuid.uuid4(),
            site_uuid=site_uuid,
            created_utc=dt.datetime.now(tz=dt.timezone.utc),
            forecast_version=forecast_version,
        )
        newly_written_rows = upsert(session, ForecastSQL, forecast.__dict__)
        written_rows.extend(newly_written_rows)

        # Get all dataframe forecast value entries for current site_uuid
        df_site: pd.DataFrame = df_forecast_values.loc[df_forecast_values["pv_uuid"] == site_uuid]

        # Filter the forecasted values by target_time
        target_times: npt.ndarray[dt.datetime] = df_site["target_datetime_utc"].unique()

        # Print a warning if there are duplicate target_times for this site's forecast
        if len(target_times) != len(df_site):
            logging.warning(
                f"duplicate target times exist in forecast {forecast.forecast_uuid} "
                f"for site {site_uuid}"
            )

        # For each target time:
        for target_time in target_times:
            # Convert the target_time string to a datetime object
            target_time_as_datetime: dt.datetime = dp.isoparse(target_time)

            datetime_interval, newly_added_rows = get_or_else_create_datetime_interval(
                session=session, start_time=target_time_as_datetime
            )
            written_rows.extend(newly_added_rows)

            # For each entry with this targettime (there should only be one)
            df_target_entries: pd.DataFrame = df_site.loc[
                df_site["target_datetime_utc"] == target_time
            ]

            # Create a ForecastValueSQL object for each forecast value, and surface as dict
            # TODO there might be quicker ways to do this, like go from pandas straight to a dict
            sql_forecast_value = ForecastValueSQL(
                forecast_uuid=forecast.forecast_uuid,
                forecast_value_uuid=uuid.uuid4(),
                datetime_interval_uuid=datetime_interval.datetime_interval_uuid,
                created_utc=dt.datetime.now(tz=dt.timezone.utc),
                forecast_generation_kw=df_target_entries.iloc[0]["forecast_kw"],
            ).__dict__
            sql_forecast_values.append(sql_forecast_value)

    # Save it to the db
    newly_added_rows = upsert(session, ForecastValueSQL, sql_forecast_values)
    written_rows.extend(newly_added_rows)

    return written_rows
