# Copyright © 2024, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE.md
import sqlite3
from typing import Union, Optional, Dict, Any

import pandas as pd
import shapely.wkb
from shapely.geometry import Polygon, Point, LineString

from polaris.network.data.data_table_cache import DataTableCache
from polaris.network.database_connection import get_srid
from polaris.network.starts_logging import logger


class DataRecord:
    """General class to table elements into Python classes

    All geometries are transformed into Shapely geometries"""

    def __init__(
        self,
        element_id: int,
        table_name: str,
        data_storage: DataTableCache,
        record: pd.DataFrame,
        conn: Optional[sqlite3.Connection] = None,
    ):
        self._data = data_storage
        self._table_name = table_name
        self.__original__: Dict[str, Any] = {}
        self.id = element_id
        self.__exists__ = True
        self.id = element_id
        self.geo: Union[Polygon, Point, LineString]
        self.__srid__ = get_srid(conn=conn, database_path=data_storage._network_file)
        self.type = ""

        if not record.shape[0]:
            data = self._data.get_table(table_name, conn)
            record = data[data.index == element_id]

        self._id_field = record.index.name
        self.__dict__[self._id_field] = element_id
        if record.shape[0]:
            self.__original__[self._id_field] = element_id
            for col in record.columns:
                val = record[col].values[0] if col != "geo" else shapely.wkb.loads(record[col].values[0])
                self.__dict__[col] = val
                self.__original__[col] = val
        else:
            self.__exists__ = False
            self.__original__[self._id_field] = None
            for col in record.columns:
                self.__dict__[col] = None
                self.__original__[col] = None
            self.id = element_id

    def save(self, conn: sqlite3.Connection):
        """Save this table record to the database"""
        if self.__exists__:
            to_change = []
            data = []
            qry = ""
            for key, val in self.__original__.items():
                if val != self.__dict__[key]:
                    if key == "geo":
                        to_change.append("geo=GeomFromWKB(?, ?)")
                        data.extend([self.geo.wkb, self.__srid__])
                    else:
                        to_change.append(f'"{key}"=?')
                        data.append(self.__dict__[key])
            if to_change:
                for key in self.__original__.keys():
                    self.__original__[key] = self.__dict__[key]
                qry = f'Update "{self._table_name}" set {",".join(to_change)} where "{self._id_field}"=?'
                data.append(self.id)
        else:
            self.__exists__ = True
            data = []
            keys = []
            string = []
            for key in self.__original__.keys():
                keys.append(f'"{key}"')
                if key == "geo":
                    data.extend([self.geo.wkb, self.__srid__])
                    string.append("GeomFromWKB(?,?)")
                else:
                    data.append(self.__dict__[key])
                    string.append("?")

            qry = f'INSERT into {self._table_name}({",".join(keys)}) VALUES({",".join(string)})'

        if data:
            try:
                conn.execute(qry, data)
                conn.commit()
            except Exception as e:
                logger.critical(qry)
                logger.critical(data)
                raise e

    def __setattr__(self, key, value):
        if value and key == "geo":
            self.__wkb__ = value.wkb
        self.__dict__[key] = value
