import random
import time
import pandas as pd
from tqdm import tqdm
from etl.find_my_epc.RetrieveFindMyEpc import RetrieveFindMyEpc
from backend.SearchEpc import SearchEpc
from utils.logger import setup_logger

logger = setup_logger()


class AssetListEpcData:

    def __init__(self, asset_list: pd.DataFrame, epc_auth_token: str):

        """
        This class handles pulling data assocaited to an asset list and performs common functions like
        getting EPC api data, retrieveing data form the find my epc website and extracting non-intrusive
        recommendations
        :param asset_list:
        """

        # Check the asset list contains the correct columns

        self.asset_list = self.check_asset_list(asset_list)
        self.epc_auth_token = epc_auth_token

        self.extracted_data = None
        self.non_invasive_recommendations = None
        self.patches = None
        self.epc_data = None

    @staticmethod
    def check_asset_list(asset_list):
        # TODO: Update this with pydantic

        return asset_list

    def get_non_invasive_recommendations(self):

        """
        Extracts non-invasive recommendations in a format that can be used by the engine
        :return:
        """

        if self.extracted_data is None:
            raise ValueError("Please run get_data first")

        self.non_invasive_recommendations = [
            {
                "uprn": r.get("uprn"),
                "address": r["address"],
                "postcode": r["postcode"],
                "recommendations": r.get("recommendations")
            } for r in self.extracted_data
        ]

    def get_patch(self):
        """

        :return:
        """
        if self.extracted_data is None:
            raise ValueError("extracted data is missing - run get_data first")

        self.patches = [
            {
                "uprn": r.get("uprn"),
                **r.get("patch")
            } for r in self.extracted_data if r.get("patch")
        ]

    def get_data(self):

        logger.info("Retrieving data for given asset list")

        # Pull the additional data
        extracted_data = []
        epc_data = []
        for _, home in tqdm(self.asset_list.iterrows(), total=len(self.asset_list)):

            add1 = home["address"]
            pc = home["postcode"]
            # Retrieve the EPC data
            epc_searcher = SearchEpc(
                address1=add1,
                postcode=pc,
                uprn=home.get("uprn"),
                auth_token=self.epc_auth_token,
                os_api_key="",
            )
            epc_searcher.ordnance_survey_client.property_type = home.get("property_type")
            epc_searcher.ordnance_survey_client.built_form = home.get("built_form")
            epc_searcher.find_property(skip_os=True)

            if epc_searcher.newest_epc is None:
                continue

            # Attempt both methods:
            try:
                find_epc_searcher = RetrieveFindMyEpc(
                    address=epc_searcher.newest_epc["address1"] + ", " + epc_searcher.newest_epc["address2"],
                    postcode=epc_searcher.newest_epc["postcode"]
                )
                find_epc_data = find_epc_searcher.retrieve_newest_find_my_epc_data()
            except Exception as e:
                logger.error(f"Error retrieving find my epc data: {e}")
                if not pd.isnull(home.get("patch")):
                    epc_searcher.newest_epc["address1"] = add1

                try:
                    find_epc_searcher = RetrieveFindMyEpc(
                        address=epc_searcher.newest_epc["address1"],
                        postcode=epc_searcher.newest_epc["postcode"]
                    )
                    find_epc_data = find_epc_searcher.retrieve_newest_find_my_epc_data()
                except Exception as e:
                    logger.error("Error retrieving find my epc data with alternative address format: {e}")
                    find_epc_data = {
                        "current_epc_rating": epc_searcher.newest_epc["current-energy-rating"],
                        "current_epc_efficiency": epc_searcher.newest_epc["current-energy-efficiency"],
                        "potential_epc_rating": None,
                        "potential_epc_efficiency": None,
                        "epc_data": {}
                    }

            # Sleep for a random amount of time between 0.5 and 1 seconds to avoid hitting the API rate limit
            time.sleep(random.sample(range(50, 100), 1)[0] / 100)

            # Every 50 requests, we sleep for 10 seconds to avoid hitting the API rate limit
            if len(extracted_data) % 50 == 0 and len(extracted_data) > 0:
                logger.info("Sleeping for 10 seconds to avoid hitting API rate limit")
                time.sleep(10)

            # We need uprn

            to_append = {
                "uprn": home.get("uprn", epc_searcher.newest_epc["uprn"]),
                "address": home["address"],
                "postcode": home["postcode"],
                **find_epc_data,
            }
            if not pd.isnull(home.get("patch")):
                to_append["patch"] = {
                    "current-energy-rating": find_epc_data["current_epc_rating"],
                    "current-energy-efficiency": find_epc_data["current_epc_efficiency"],
                    "potential-energy-rating": find_epc_data["potential_epc_rating"],
                    "potential-energy-efficiency": find_epc_data["potential_epc_efficiency"],
                    **find_epc_data["epc_data"]
                }

            extracted_data.append(to_append)
            epc_data.append(epc_searcher.newest_epc)

        self.extracted_data = extracted_data
        self.epc_data = epc_data
        logger.info("Data Extrction complete")
