from django.core.management.base import BaseCommand
from django.apps import apps
from django.db import connection, transaction


class Command(BaseCommand):
    help = "Synchronizes the geojson field (and lon/lat if missing) from geom for all assets/models using PostGIS SQL."

    def add_arguments(self, parser):
        parser.add_argument(
            '--dry-run',
            action='store_true',
            help='Perform a dry run without modifying the database.',
        )
        parser.add_argument(
            '--only-missing',
            action='store_true',
            help='Only update records where geojson is null or empty.',
        )

    def handle(self, *args, **options):
        dry_run = options['dry_run']
        only_missing = options['only_missing']

        total_scanned = 0
        total_updated = 0

        with connection.cursor() as cursor:
            for model in apps.get_models():
                fields = [f.name for f in model._meta.get_fields()]
                if 'geom' not in fields or 'geojson' not in fields:
                    continue

                table = model._meta.db_table
                is_json_field = model._meta.get_field('geojson').get_internal_type() == 'JSONField'
                has_lon = 'lon' in fields
                has_lat = 'lat' in fields

                # Count total candidates with non-null geom
                cursor.execute(f'SELECT COUNT(*) FROM "{table}" WHERE geom IS NOT NULL;')
                count = cursor.fetchone()[0]
                if count == 0:
                    continue

                self.stdout.write(f"Processing {model._meta.app_label}.{model.__name__} ({table}, {count} candidate records)...")

                if only_missing:
                    where_clause = "geom IS NOT NULL AND (geojson IS NULL OR geojson::text = '' OR geojson::text = 'null')"
                else:
                    where_clause = "geom IS NOT NULL"

                cursor.execute(f'SELECT COUNT(*) FROM "{table}" WHERE {where_clause};')
                records_to_update = cursor.fetchone()[0]

                if records_to_update == 0:
                    self.stdout.write(self.style.SUCCESS(f"  {model.__name__}: 0 records need update."))
                    continue

                if dry_run:
                    self.stdout.write(self.style.SUCCESS(f"  [DRY-RUN] {model.__name__}: Would update {records_to_update}/{count} records."))
                    total_scanned += count
                    total_updated += records_to_update
                    continue

                cast_geojson = "ST_AsGeoJSON(ST_Transform(geom, 4326))::jsonb" if is_json_field else "ST_AsGeoJSON(ST_Transform(geom, 4326))"

                set_clauses = [f"geojson = {cast_geojson}"]
                if has_lon:
                    set_clauses.append("lon = COALESCE(lon, ST_X(ST_PointOnSurface(ST_Transform(geom, 4326))))")
                if has_lat:
                    set_clauses.append("lat = COALESCE(lat, ST_Y(ST_PointOnSurface(ST_Transform(geom, 4326))))")

                set_str = ", ".join(set_clauses)
                update_sql = f'UPDATE "{table}" SET {set_str} WHERE {where_clause};'

                with transaction.atomic():
                    cursor.execute(update_sql)
                    updated_count = cursor.rowcount

                self.stdout.write(self.style.SUCCESS(f"  {model.__name__}: {updated_count}/{count} records updated in DB."))
                total_scanned += count
                total_updated += updated_count

        prefix = "[DRY-RUN] Would update" if dry_run else "Successfully updated"
        self.stdout.write(self.style.SUCCESS(f"\n{prefix} {total_updated} / {total_scanned} records across all models."))