from django.core.management.base import BaseCommand
from django.db import connection

class Command(BaseCommand):
    help = 'Fix location fields in parks and historical records'

    def handle(self, *args, **options):
        with connection.cursor() as cursor:
            # Check if Unknown country exists
            cursor.execute("""
                SELECT id FROM cities_light_country WHERE name = 'Unknown' LIMIT 1;
            """)
            result = cursor.fetchone()
            if result:
                default_country_id = result[0]
            else:
                cursor.execute("""
                    INSERT INTO cities_light_country (name, name_ascii, slug, geoname_id, alternate_names)
                    VALUES ('Unknown', 'Unknown', 'unknown', 0, '')
                    RETURNING id;
                """)
                default_country_id = cursor.fetchone()[0]

            # Check if Unknown region exists
            cursor.execute("""
                SELECT id FROM cities_light_region 
                WHERE name = 'Unknown' AND country_id = %s LIMIT 1;
            """, [default_country_id])
            result = cursor.fetchone()
            if result:
                default_region_id = result[0]
            else:
                cursor.execute("""
                    INSERT INTO cities_light_region (name, name_ascii, slug, geoname_id, alternate_names, country_id, display_name)
                    VALUES ('Unknown', 'Unknown', 'unknown', 0, '', %s, 'Unknown')
                    RETURNING id;
                """, [default_country_id])
                default_region_id = cursor.fetchone()[0]

            # Check if Unknown city exists
            cursor.execute("""
                SELECT id FROM cities_light_city 
                WHERE name = 'Unknown' AND region_id = %s LIMIT 1;
            """, [default_region_id])
            result = cursor.fetchone()
            if result:
                default_city_id = result[0]
            else:
                cursor.execute("""
                    INSERT INTO cities_light_city (
                        name, name_ascii, slug, geoname_id, alternate_names, 
                        region_id, country_id, display_name,
                        latitude, longitude, population
                    )
                    VALUES (
                        'Unknown', 'Unknown', 'unknown', 0, '', 
                        %s, %s, 'Unknown',
                        0, 0, 0
                    )
                    RETURNING id;
                """, [default_region_id, default_country_id])
                default_city_id = cursor.fetchone()[0]

            # Update parks with null locations
            cursor.execute("""
                UPDATE parks_park
                SET country_id = %s,
                    region_id = %s,
                    city_id = %s,
                    location = 'Unknown, Unknown, Unknown'
                WHERE country_id IS NULL
                   OR region_id IS NULL
                   OR city_id IS NULL
                   OR location IS NULL;
            """, [default_country_id, default_region_id, default_city_id])

            # Update historical records with null locations
            cursor.execute("""
                UPDATE parks_historicalpark
                SET country_id = %s,
                    region_id = %s,
                    city_id = %s,
                    location = 'Unknown, Unknown, Unknown'
                WHERE country_id IS NULL
                   OR region_id IS NULL
                   OR city_id IS NULL
                   OR location IS NULL;
            """, [default_country_id, default_region_id, default_city_id])

            self.stdout.write(self.style.SUCCESS('Successfully fixed location fields'))