# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
import numpy as np
import geopandas as gpd
from qgis.core import QgsFeature, QgsField, QgsProject, QgsVectorLayer, QgsGeometry
from qgis.PyQt.QtCore import QVariant


def layer_from_geodataframe(gdf: gpd.GeoDataFrame, geo_type: str, layer_name: str, add_layer=True) -> QgsVectorLayer:
    # create layer
    vl = QgsVectorLayer(f"{geo_type}?crs=epsg:{gdf.crs.to_epsg()}", layer_name, "memory")
    pr = vl.dataProvider()

    if "geometry" not in gdf:
        gdf.rename_geometry("geometry", inplace=True)

    # add fields
    def qgs_type(ftype):
        return QVariant.Double if "float" in ftype.name else QVariant.Int if "int" in ftype.name else QVariant.String

    field_names = [x for x in gdf.columns if x != "geometry"]

    types = [qgs_type(gdf.dtypes[fname]) for fname in field_names]
    attributes = [QgsField(fname, dtype) for fname, dtype in zip(field_names, types)]
    pr.addAttributes(attributes)
    vl.updateFields()  # tell the vector layer to fetch changes from the provider

    # Add records
    features = []
    gdf.replace({np.nan: None}, inplace=True)
    for _, record in gdf.iterrows():
        fet = QgsFeature()
        fet.setAttributes(record[field_names].to_list())
        fet.setGeometry(QgsGeometry.fromWkt(record.geometry.wkt))
        features.append(fet)
    pr.addFeatures(features)
    vl.updateExtents()

    if add_layer:
        QgsProject.instance().addMapLayer(vl)

    # returns the layer handle
    return vl
