import os
import json
import shutil
import tempfile

from qgis.core import (
    QgsProject,
    QgsVectorFileWriter,
    QgsVectorLayer,
    QgsCoordinateTransformContext,
    QgsFeature,
)


SESSION_FILE = "session.json"
SESSION_GPKG = "session.gpkg"
SOURCE_DIR = "source_zip"


def _ensure_dir(path):
    os.makedirs(path, exist_ok=True)


def _copy_source_zip_contents(package, session_dir):
    source_dir = os.path.join(session_dir, SOURCE_DIR)
    _ensure_dir(source_dir)

    for name, path in package.files.items():
        target = os.path.join(source_dir, name)
        shutil.copy2(path, target)

    return source_dir


def _write_session_json(session_dir, data):
    session_json_path = os.path.join(session_dir, SESSION_FILE)
    with open(session_json_path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    return session_json_path


def _read_session_json(session_dir):
    session_json_path = os.path.join(session_dir, SESSION_FILE)
    if not os.path.exists(session_json_path):
        raise Exception("session.json not found in selected session folder.")
    with open(session_json_path, "r", encoding="utf-8") as f:
        return json.load(f)


def _open_gpkg_layer(gpkg_path, layer_name):
    uri = f"{gpkg_path}|layername={layer_name}"
    layer = QgsVectorLayer(uri, layer_name, "ogr")
    if not layer.isValid():
        raise Exception(f"Could not open layer {layer_name} from session.gpkg")
    return layer


def _save_layer_to_gpkg(layer, gpkg_path, layer_name, create_new_file=False):
    options = QgsVectorFileWriter.SaveVectorOptions()
    options.driverName = "GPKG"
    options.layerName = layer_name
    options.fileEncoding = "UTF-8"

    if create_new_file:
        options.actionOnExistingFile = QgsVectorFileWriter.CreateOrOverwriteFile
    else:
        options.actionOnExistingFile = QgsVectorFileWriter.CreateOrOverwriteLayer

    error_code, error_message, new_file_name, new_layer = QgsVectorFileWriter.writeAsVectorFormatV3(
        layer,
        gpkg_path,
        QgsCoordinateTransformContext(),
        options
    )

    if error_code != QgsVectorFileWriter.NoError:
        raise Exception(f"Could not save {layer_name} to session.gpkg: {error_message}")


def _clone_layer_to_memory(layer, clone_name):
    """
    Krijon një memory clone të plotë të layer-it:
    - fields
    - features
    - geometry
    - CRS
    """
    crs_authid = layer.crs().authid() if layer.crs().isValid() else "EPSG:6316"
    geom_type = layer.wkbType()
    geom_name = layer.wkbType()

    # përdorim provider memory me Polygon, sepse layer-at tanë janë polygon
    mem_uri = f"Polygon?crs={crs_authid}"
    mem_layer = QgsVectorLayer(mem_uri, clone_name, "memory")
    if not mem_layer.isValid():
        raise Exception(f"Could not create memory clone for layer {clone_name}")

    pr = mem_layer.dataProvider()
    pr.addAttributes(layer.fields())
    mem_layer.updateFields()

    new_features = []
    for feat in layer.getFeatures():
        new_feat = QgsFeature(mem_layer.fields())
        new_feat.setGeometry(feat.geometry())
        new_feat.setAttributes(feat.attributes())
        new_features.append(new_feat)

    ok, _ = pr.addFeatures(new_features)
    if not ok:
        raise Exception(f"Could not copy features into memory clone: {clone_name}")

    mem_layer.updateExtents()
    return mem_layer


def create_session_from_layers(package, parcels_layer, buildings_layer, session_dir):
    if package is None:
        raise Exception("No source package loaded.")

    if parcels_layer is None or buildings_layer is None:
        raise Exception("Layers are missing.")

    _ensure_dir(session_dir)

    gpkg_path = os.path.join(session_dir, SESSION_GPKG)
    source_dir = _copy_source_zip_contents(package, session_dir)

    _save_layer_to_gpkg(parcels_layer, gpkg_path, "PARC_TRANS", create_new_file=True)
    _save_layer_to_gpkg(buildings_layer, gpkg_path, "BUILDINGS_TRANS", create_new_file=False)

    session_data = {
        "session_version": 1,
        "gpkg": SESSION_GPKG,
        "source_dir": SOURCE_DIR,
        "layers": {
            "parcels": "PARC_TRANS",
            "buildings": "BUILDINGS_TRANS",
        }
    }
    _write_session_json(session_dir, session_data)

    return {
        "session_dir": session_dir,
        "gpkg_path": gpkg_path,
        "source_dir": source_dir,
        "parcels_layer": _open_gpkg_layer(gpkg_path, "PARC_TRANS"),
        "buildings_layer": _open_gpkg_layer(gpkg_path, "BUILDINGS_TRANS"),
    }


def open_existing_session(session_dir):
    data = _read_session_json(session_dir)

    gpkg_path = os.path.join(session_dir, data["gpkg"])
    source_dir = os.path.join(session_dir, data["source_dir"])

    if not os.path.exists(gpkg_path):
        raise Exception("session.gpkg not found.")

    if not os.path.isdir(source_dir):
        raise Exception("source_zip folder not found.")

    parcels_layer_name = data["layers"]["parcels"]
    buildings_layer_name = data["layers"]["buildings"]

    parcels_layer = _open_gpkg_layer(gpkg_path, parcels_layer_name)
    buildings_layer = _open_gpkg_layer(gpkg_path, buildings_layer_name)

    files = {}
    for name in ["parcels_o.gml", "parcels.gml", "buildings_o.gml", "buildings.gml", "elab.xml"]:
        path = os.path.join(source_dir, name)
        if os.path.exists(path):
            files[name] = path

    if len(files) < 5:
        raise Exception("Session source_zip folder is incomplete.")

    return {
        "session_dir": session_dir,
        "gpkg_path": gpkg_path,
        "source_dir": source_dir,
        "parcels_layer": parcels_layer,
        "buildings_layer": buildings_layer,
        "files": files,
    }


def save_session_layers(session_info):
    if session_info is None:
        raise Exception("No session available.")

    gpkg_path = session_info["gpkg_path"]
    session_dir = session_info["session_dir"]
    parcels_layer = session_info["parcels_layer"]
    buildings_layer = session_info["buildings_layer"]

    if parcels_layer is None or buildings_layer is None:
        raise Exception("Session layers missing.")

    if not gpkg_path or not session_dir:
        raise Exception("Session paths are missing.")

    project = QgsProject.instance()

    # 1. Krijo memory clones që s’kanë lock me gpkg
    parcels_mem = _clone_layer_to_memory(parcels_layer, "PARC_TRANS")
    buildings_mem = _clone_layer_to_memory(buildings_layer, "BUILDINGS_TRANS")

    # 2. Hiq layer-at aktualë nga projekti për të liruar çdo lock të mundshëm
    if project.mapLayer(parcels_layer.id()):
        project.removeMapLayer(parcels_layer.id())
    if project.mapLayer(buildings_layer.id()):
        project.removeMapLayer(buildings_layer.id())

    # 3. Shkruaj në një GPKG të përkohshëm
    tmp_dir = tempfile.mkdtemp(prefix="cadastre_save_")
    tmp_gpkg = os.path.join(tmp_dir, SESSION_GPKG)

    _save_layer_to_gpkg(parcels_mem, tmp_gpkg, "PARC_TRANS", create_new_file=True)
    _save_layer_to_gpkg(buildings_mem, tmp_gpkg, "BUILDINGS_TRANS", create_new_file=False)

    # 4. Backup opsional i versionit të vjetër
    backup_gpkg = gpkg_path + ".bak"
    if os.path.exists(backup_gpkg):
        try:
            os.remove(backup_gpkg)
        except Exception:
            pass

    if os.path.exists(gpkg_path):
        os.replace(gpkg_path, backup_gpkg)

    # 5. Zëvendëso session.gpkg me versionin e ri
    shutil.copy2(tmp_gpkg, gpkg_path)

    # 6. Rihap layer-at nga session.gpkg i ri
    new_parcels_layer = _open_gpkg_layer(gpkg_path, "PARC_TRANS")
    new_buildings_layer = _open_gpkg_layer(gpkg_path, "BUILDINGS_TRANS")

    project.addMapLayer(new_parcels_layer)
    project.addMapLayer(new_buildings_layer)

    session_info["parcels_layer"] = new_parcels_layer
    session_info["buildings_layer"] = new_buildings_layer

    return gpkg_path