from django.core.management.base import BaseCommand

from ._spreadsheet_constants import *
from right_tree.api.models import Habitat, Zone


def load_habitat_zone_data(sheet):
    """ Loads habitat and zone objects from data defined in the spreadsheet.
    """

    habitats = set()
    zones = set()

    current_habitat = current_zone = current_variant = current_refined_variant = None

    for col in sheet.iter_cols(min_col=SITE_DATA_START_COL, max_col=SITE_DATA_STOP_COL, min_row=HABITAT_ROW, max_row=ZONE_REFINED_VARIANT_ROW):
        habitat, zone_name, zone_variant, zone_refined_variant = (cell.value for cell in col)

        if habitat is not None:
            current_habitat = habitat
            current_zone = current_variant = current_refined_variant = None
        
        if zone_name is not None:
            current_zone = zone_name
            current_variant = current_refined_variant = None

        if zone_variant is not None:
            current_variant = zone_variant
            current_refined_variant = None

        current_refined_variant = zone_refined_variant if zone_refined_variant is not None else current_refined_variant

        habitats.add(current_habitat)
        zones.add((current_habitat, current_zone, current_variant, current_refined_variant, col[0].column_letter))

    # Create habitats - generally this is not needed for updating
    # for habitat in habitats:
    #     habitat_obj = Habitat(name=habitat)
    #     habitat_obj.save()

    for habitat, zone, variant, refined_variant, col_letter in zones:
        
        # Retrieve existing objects
        habitat_obj = Habitat.objects.get(name=habitat)
        zone_obj = Zone.objects.get(related_svg_segment=col_letter)
        
        # Update fields if required
        zone_obj.name = zone
        zone_obj.variant = variant
        zone_obj.refined_variant = refined_variant
        zone_obj.habitat = habitat_obj

        # Save the object
        zone_obj.save()
class Command(BaseCommand):
    help = 'Ingests the site spreadsheet data into the database'

    def handle(self, *args, **options):
        load_habitat_zone_data(SPREADSHEET)