import datetime
import glob
import logging
import os
import random
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import geopandas as gpd
import osmnx as ox
import pandas as pd
import requests
from PIL import Image
from shapely.errors import ShapelyDeprecationWarning
from tqdm import tqdm

import zensvi.download.kartaview.download_functions as kv
from zensvi.download.base import BaseDownloader
from zensvi.download.utils.helpers import check_and_buffer, standardize_column_names
from zensvi.utils.log import Logger

warnings.filterwarnings("ignore", category=ShapelyDeprecationWarning)
logging.getLogger("mapillary.utils.client").setLevel(logging.WARNING)


class KVDownloader(BaseDownloader):
    """KartaView Downloader class.

    Args:
        log_path (str, optional): Path to the log file. Defaults to None.
        max_workers (int, optional): Number of workers for parallel processing. Defaults to None.
    """

    def __init__(self, log_path=None, max_workers=None):
        super().__init__(log_path)
        self._max_workers = max_workers
        # initialize the logger
        if log_path is not None:
            # make log directory
            Path(log_path).parent.mkdir(parents=True, exist_ok=True)
            self.logger = Logger(log_path)
        else:
            self.logger = None

    @property
    def max_workers(self):
        """Property for the number of workers for parallel processing.

        Returns:
            int: max_workers
        """
        return self._max_workers

    @max_workers.setter
    def max_workers(self, max_workers):
        if max_workers is None:
            self._max_workers = min(32, os.cpu_count() + 4)
        else:
            self._max_workers = max_workers

    def _read_pids(self, path_pid):
        pid_df = pd.read_csv(path_pid)
        # drop NA values in id columns
        pid_df = pid_df.dropna(subset=["id"])
        # get unique pids (ie "id" columns) as a list
        pids = pid_df["id"].astype("int64").unique().tolist()
        return pids

    def _set_dirs(self, dir_output):
        # set dir_output as attribute and create the directory
        self.dir_output = Path(dir_output)
        self.dir_output.mkdir(parents=True, exist_ok=True)
        self.pids_url = self.dir_output / "pids_urls_kv.csv"
        # # set dir_cache as attribute and create the directory
        # self.dir_cache = self.dir_output / "cache_zensvi_kv"
        # self.dir_cache.mkdir(parents=True, exist_ok=True)
        # # set other cache directories
        # self.cache_lat_lon = self.dir_cache / "lat_lon_kv.csv"
        # self.cache_pids_raw = self.dir_cache / "pids_raw_kv.csv"

    def _get_pids_from_gdf(self, gdf):
        # set crs to EPSG:4326 if it's None
        if gdf.crs is None:
            gdf = gdf.set_crs("EPSG:4326")
        elif gdf.crs != "EPSG:4326":
            # convert to EPSG:4326
            gdf = gdf.to_crs("EPSG:4326")
        # use get pids with kv.get_points_in_shape
        result_df = kv.get_points_in_shape(gdf)
        return result_df

    def _get_raw_pids(self, **kwargs):
        # if self.cache_pids_raw.exists():
        #     pid = pd.read_csv(self.cache_pids_raw)
        #     print("The raw image IDs have been read from the cache")
        #     return pid

        # input: lat and lon
        if kwargs["lat"] is not None and kwargs["lon"] is not None:
            # create a geodataframe with lat and lon
            df = pd.DataFrame({"lat": [kwargs["lat"]], "lon": [kwargs["lon"]]})
            gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.lon, df.lat), crs="EPSG:4326")

        # input: csv file
        elif kwargs["input_csv_file"] != "":
            df = pd.read_csv(kwargs["input_csv_file"])
            df = standardize_column_names(df)
            gdf = gpd.GeoDataFrame(
                df,
                geometry=gpd.points_from_xy(df.longitude, df.latitude),
                crs="EPSG:4326",
            )

        # input: shapefile
        elif kwargs["input_shp_file"] != "":
            gdf = gpd.read_file(kwargs["input_shp_file"])

        # input: place name
        elif kwargs["input_place_name"] != "":
            print("Geocoding the input place name")
            gdf = ox.geocoder.geocode_to_gdf(kwargs["input_place_name"])
            # raise error if the input_place_name is not found
            if len(gdf) == 0:
                raise ValueError("The input_place_name is not found. Please try another place name.")
        else:
            raise ValueError("Please input the lat and lon, csv file, or shapefile.")

        # check geometry type and buffer
        gdf = check_and_buffer(gdf, kwargs["buffer"])
        # get pid
        pid = self._get_pids_from_gdf(gdf)
        pid = self._filter_pids_date(pid, kwargs["start_date"], kwargs["end_date"])

        return pid

    def _filter_pids_date(self, pid_df, start_date, end_date):
        # create a temporary column date from captured_at (milliseconds from Unix epoch)
        pid_df["date"] = pd.to_datetime(pid_df["shotDate"], format="%Y-%m-%d %H:%M:%S")
        # check if start_date and end_date are in the correct format with regex. If not, raise error
        if start_date is not None:
            try:
                start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d")
            except ValueError:
                raise ValueError("Incorrect start_date format, should be YYYY-MM-DD")
        if end_date is not None:
            try:
                end_date = datetime.datetime.strptime(end_date, "%Y-%m-%d")
            except ValueError:
                raise ValueError("Incorrect end_date format, should be YYYY-MM-DD")
        # if start_date is not None, filter out the rows with date < start_date
        pid_df = pid_df[pid_df["date"] >= start_date] if start_date is not None else pid_df
        # if end_date is not None, filter out the rows with date > end_date
        pid_df = pid_df[pid_df["date"] <= end_date] if end_date is not None else pid_df
        # drop the temporary column date
        pid_df = pid_df.drop(columns="date")
        return pid_df

    def _get_pids(self, path_pid, **kwargs):
        # get raw pid
        pid = self._get_raw_pids(**kwargs)

        if pid is None:
            print("There is no image ID to download")
            return

        # move the "id" column to the first column
        pid = pid[["id"] + [col for col in pid.columns if col != "id"]]

        pid.to_csv(path_pid, index=False)
        print("The image IDs have been saved to {}".format(path_pid))

    def _get_urls_kv(self, path_pid):

        pid = pd.read_csv(path_pid)
        if len(pid) == 0:
            print("There is no image ID to download")
            return
        urls = pid[["id", "fileurlProc"]].rename(columns={"fileurlProc": "url"})
        urls.to_csv(self.pids_url, index=False)

    def _download_images_kv(self, path_pid, cropped, batch_size, start_date, end_date):
        checkpoints = glob.glob(str(self.panorama_output / "**/*.png"), recursive=True)

        # Read already downloaded images and convert to ids
        downloaded_ids = set([int(Path(file_path).stem) for file_path in checkpoints])  # Use set for faster operations

        pid_df = pd.read_csv(path_pid).dropna(subset=["id"])
        pid_df["id"] = pid_df["id"].astype("int64")
        urls_df = pd.read_csv(self.pids_url)
        urls_df["id"] = urls_df["id"].astype("int64")
        # merge pid_df and urls_df
        urls_df = urls_df.merge(pid_df, on="id", how="left")
        # filter out the rows by date
        urls_df = self._filter_pids_date(urls_df, start_date, end_date)

        # Filter out the ids that have already been processed
        urls_df = urls_df[~urls_df["id"].isin(downloaded_ids)]  # Use isin for efficient operation

        def worker(row, output_dir, cropped, max_retries=5):
            url, panoid = row.url, row.id
            user_agent = random.choice(self._user_agents)
            proxy = random.choice(self._proxies)

            image_name = f"{panoid}.png"  # Use id for file name
            image_path = output_dir / image_name

            for retry in range(max_retries):
                try:
                    response = requests.get(url, headers=user_agent, proxies=proxy, timeout=10)
                    if response.status_code == 200:
                        with open(image_path, "wb") as f:
                            f.write(response.content)

                        if cropped:
                            img = Image.open(image_path)
                            w, h = img.size
                            img_cropped = img.crop((0, 0, w, h // 2))
                            img_cropped.save(image_path)
                        break
                    else:
                        if retry == max_retries - 1:  # Last retry
                            if self.logger is not None:
                                self.logger.log_failed_pids(panoid)
                except Exception as e:
                    if retry == max_retries - 1:  # Last retry
                        if self.logger is not None:
                            self.logger.log_failed_pids(panoid)
                        print(f"Error: {e}")
                    else:
                        print(f"Retry {retry + 1}/{max_retries} failed: {e}")

        num_batches = (len(urls_df) + batch_size - 1) // batch_size

        # Calculate current highest batch number
        existing_batches = glob.glob(str(self.panorama_output / "batch_*"))
        existing_batch_numbers = [int(Path(batch).name.split("_")[-1]) for batch in existing_batches]
        start_batch_number = max(existing_batch_numbers, default=0)

        for i in tqdm(
            range(start_batch_number, start_batch_number + num_batches),
            desc=f"Downloading images by batch size {min(batch_size, len(urls_df))}",
        ):
            # Create a new sub-folder for each batch
            batch_out_path = self.panorama_output / f"batch_{i+1}"
            batch_out_path.mkdir(exist_ok=True)

            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                batch_futures = {
                    executor.submit(worker, row, batch_out_path, cropped): row.id
                    for row in urls_df.iloc[
                        (i - start_batch_number) * batch_size : (i + 1 - start_batch_number) * batch_size
                    ].itertuples()
                }
                for future in tqdm(
                    as_completed(batch_futures),
                    total=len(batch_futures),
                    desc=f"Downloading images for batch #{i+1}",
                ):
                    try:
                        future.result()
                    except Exception as e:
                        print(f"Error: {e}")

    def download_svi(
        self,
        dir_output,
        path_pid=None,
        lat=None,
        lon=None,
        input_csv_file="",
        input_shp_file="",
        input_place_name="",
        buffer=10,
        update_pids=False,
        cropped=False,
        batch_size=1000,
        start_date=None,
        end_date=None,
        metadata_only=False,
        max_workers=None,
    ):
        """Downloads street view images from KartaView using specified parameters.

        Args:
            dir_output (str): Directory where output files and images will be stored.
            path_pid (str, optional): Path to a file containing panorama IDs. If not provided, IDs will be fetched based on other parameters.
            lat (float, optional): Latitude to fetch panorama IDs around this point. Must be used with `lon`.
            lon (float, optional): Longitude to fetch panorama IDs around this point. Must be used with `lat`.
            input_csv_file (str, optional): Path to a CSV file containing locations for which to fetch panorama IDs.
            input_shp_file (str, optional): Path to a shapefile containing geographic locations for fetching panorama IDs.
            input_place_name (str, optional): A place name for geocoding to fetch panorama IDs.
            buffer (int, optional): Buffer size in meters to expand the geographic area for panorama ID fetching.
            update_pids (bool, optional): If True, will update panorama IDs even if a valid `path_pid` is provided. Defaults to False.
            cropped (bool, optional): If True, images will be cropped to the upper half. Defaults to False.
            batch_size (int, optional): Number of images to process in each batch. Defaults to 1000.
            start_date (str, optional): Start date (YYYY-MM-DD) to filter images by capture date.
            end_date (str, optional): End date (YYYY-MM-DD) to filter images by capture date.
            metadata_only (bool, optional): If True, skips downloading images and only fetches metadata. Defaults to False.
            max_workers (int, optional): Number of workers for parallel processing. If None, it will be set to min(32, os.cpu_count() + 4).

        Returns:
            None: This method does not return a value but will save files directly to the specified output directory.

        Raises:
            ValueError: If required parameters for fetching panorama IDs are not adequately specified.
            FileNotFoundError: If `path_pid` is specified but the file does not exist.

        Notes:
            This method logs significant events and errors, making it suitable for both interactive usage and automated workflows.
        """
        if self.logger is not None:
            # record the arguments
            self.logger.log_args(
                "KVDownloader download_svi",
                dir_output=dir_output,
                path_pid=path_pid,
                lat=lat,
                lon=lon,
                input_csv_file=input_csv_file,
                input_shp_file=input_shp_file,
                input_place_name=input_place_name,
                buffer=buffer,
                update_pids=update_pids,
                cropped=cropped,
                batch_size=batch_size,
                start_date=start_date,
                end_date=end_date,
                metadata_only=metadata_only,
                max_workers=max_workers,
            )
        # Set max_workers
        self.max_workers = max_workers

        # set necessary directories
        self._set_dirs(dir_output)

        # call _get_pids function first if path_pid is None
        if path_pid is None:
            print("Getting pids...")
            path_pid = self.dir_output / "kv_pids.csv"
            if path_pid.exists() & (not update_pids):
                print("update_pids is set to False. So the following csv file will be used: {}".format(path_pid))
            else:
                self._get_pids(
                    path_pid,
                    lat=lat,
                    lon=lon,
                    input_csv_file=input_csv_file,
                    input_shp_file=input_shp_file,
                    input_place_name=input_place_name,
                    buffer=buffer,
                    start_date=start_date,
                    end_date=end_date,
                )
        else:
            # check if the path_pid exists
            if Path(path_pid).exists():
                print("The following csv file will be used: {}".format(path_pid))
            else:
                self._get_pids(
                    path_pid,
                    lat=lat,
                    lon=lon,
                    input_csv_file=input_csv_file,
                    input_shp_file=input_shp_file,
                    input_place_name=input_place_name,
                    buffer=buffer,
                    start_date=start_date,
                    end_date=end_date,
                )
        # stop if metadata_only is True
        if metadata_only:
            print("The metadata has been downloaded")
            return

        # create a folder within self.dir_output
        self.panorama_output = self.dir_output / "kv_svi"
        self.panorama_output.mkdir(parents=True, exist_ok=True)

        # get urls
        if path_pid.exists():
            self._get_urls_kv(path_pid)
            # download images
            self._download_images_kv(path_pid, cropped, batch_size, start_date, end_date)
        else:
            print("There is no imagae ID to download within the given input parameters")

        # # delete the cache directory
        # if self.dir_cache.exists():
        #     shutil.rmtree(self.dir_cache)
        #     print("The cache directory has been deleted")
