import React, { useState, useEffect } from 'react';
import {
  Box,
  TextField,
  Autocomplete,
  Typography,
  Grid,
  Chip,
  FormControlLabel,
  Checkbox
} from '@mui/material';
import { LocationOn, Public, Home } from '@mui/icons-material';
import { Country, State, City } from 'country-state-city';
import type { ICountry, IState, ICity } from 'country-state-city';
// Import from your types file - adjust path as needed
import type { Location } from 'types/types';

interface LocationInputProps {
  value?: Partial<Location>;
  onChange: (location: Partial<Location>) => void;
  error?: boolean;
  helperText?: string;
  required?: boolean;
  disabled?: boolean;
  showCity?: boolean;
}

const LocationInput: React.FC<LocationInputProps> = ({
  value = {},
  onChange,
  error = false,
  helperText,
  required = false,
  disabled = false,
  showCity = false
}) => {
  // Get all countries from the library
  const allCountries = Country.getAllCountries();
  
  const [selectedCountry, setSelectedCountry] = useState<ICountry | null>(
    value.country ? allCountries.find(c => c.name === value.country) || null : null
  );
  const [selectedState, setSelectedState] = useState<IState | null>(null);
  const [selectedCity, setSelectedCity] = useState<ICity | null>(null);
  const [isRemote, setIsRemote] = useState<boolean>(value.remote || false);

  // Get states for selected country
  const availableStates = selectedCountry ? State.getStatesOfCountry(selectedCountry.isoCode) : [];
  
  // Get cities for selected state
  const availableCities = selectedCountry && selectedState 
    ? City.getCitiesOfState(selectedCountry.isoCode, selectedState.isoCode) 
    : [];

  // Initialize state and city from value prop
  useEffect(() => {
    if (selectedCountry && value.state) {
      const stateMatch = availableStates.find(s => s.name === value.state);
      setSelectedState(stateMatch || null);
    }
  }, [selectedCountry, value.state, availableStates]);

  useEffect(() => {
    if (selectedCountry && selectedState && value.city && showCity) {
      const cityMatch = availableCities.find(c => c.name === value.city);
      setSelectedCity(cityMatch || null);
    }
  }, [selectedCountry, selectedState, value.city, availableCities, showCity]);

  // Update parent component when values change
  useEffect(() => {
    const newLocation: Partial<Location> = {};
    
    if (selectedCountry) {
      newLocation.country = selectedCountry.name;
    }
    
    if (selectedState) {
      newLocation.state = selectedState.name;
    }
    
    if (selectedCity && showCity) {
      newLocation.city = selectedCity.name;
    }
    
    if (isRemote) {
      newLocation.remote = isRemote;
    }
    
    // Only call onChange if there's actual data or if clearing
    if (Object.keys(newLocation).length > 0 || (value.country || value.state || value.city)) {
      onChange(newLocation);
    }
  }, [selectedCountry, selectedState, selectedCity, isRemote, onChange, value.country, value.state, value.city, showCity]);

  const handleCountryChange = (event: any, newValue: ICountry | null) => {
    setSelectedCountry(newValue);
    // Clear state and city when country changes
    setSelectedState(null);
    setSelectedCity(null);
  };

  const handleStateChange = (event: any, newValue: IState | null) => {
    setSelectedState(newValue);
    // Clear city when state changes
    setSelectedCity(null);
  };

  const handleCityChange = (event: any, newValue: ICity | null) => {
    setSelectedCity(newValue);
  };

  const handleRemoteToggle = (event: React.ChangeEvent<HTMLInputElement>) => {
    setIsRemote(event.target.checked);
  };

  return (
    <Box>
      <Typography variant="h6" gutterBottom sx={{ display: 'flex', alignItems: 'center', gap: 1 }}>
        <LocationOn color="primary" />
        Location {required && <span style={{ color: 'red' }}>*</span>}
      </Typography>
      
      <Grid container spacing={2}>
        {/* Country Selection */}
        <Grid size={{ xs: 12, sm: showCity ? 4 : 6 }}>
          <Autocomplete
            value={selectedCountry}
            onChange={handleCountryChange}
            options={allCountries}
            getOptionLabel={(option) => option.name}
            disabled={disabled}
            renderInput={(params) => (
              <TextField
                {...params}
                label="Country"
                variant="outlined"
                required={required}
                error={error && required && !selectedCountry}
                helperText={error && required && !selectedCountry ? 'Country is required' : helperText}
                InputProps={{
                  ...params.InputProps,
                  startAdornment: <Public sx={{ mr: 1, color: 'text.secondary' }} />
                }}
              />
            )}
            renderOption={(props, option) => (
              <Box component="li" {...props} key={option.isoCode}>
                <img
                  loading="lazy"
                  width="20"
                  src={`https://flagcdn.com/w20/${option.isoCode.toLowerCase()}.png`}
                  srcSet={`https://flagcdn.com/w40/${option.isoCode.toLowerCase()}.png 2x`}
                  alt=""
                  style={{ marginRight: 8 }}
                />
                {option.name}
              </Box>
            )}
          />
        </Grid>

        {/* State/Region Selection */}
        {selectedCountry && (
          <Grid size={{ xs: 12, sm: showCity ? 4 : 6 }}>
            <Autocomplete
              value={selectedState}
              onChange={handleStateChange}
              options={availableStates}
              getOptionLabel={(option) => option.name}
              disabled={disabled || availableStates.length === 0}
              renderInput={(params) => (
                <TextField
                  {...params}
                  label="State/Region"
                  variant="outlined"
                  placeholder={availableStates.length > 0 ? "Select state/region" : "No states available"}
                />
              )}
            />
          </Grid>
        )}

        {/* City Selection */}
        {showCity && selectedCountry && selectedState && (
          <Grid size={{ xs: 12, sm: 4 }}>
            <Autocomplete
              value={selectedCity}
              onChange={handleCityChange}
              options={availableCities}
              getOptionLabel={(option) => option.name}
              disabled={disabled || availableCities.length === 0}
              renderInput={(params) => (
                <TextField
                  {...params}
                  label="City"
                  variant="outlined"
                  placeholder={availableCities.length > 0 ? "Select city" : "No cities available"}
                  InputProps={{
                    ...params.InputProps,
                    startAdornment: <Home sx={{ mr: 1, color: 'text.secondary' }} />
                  }}
                />
              )}
            />
          </Grid>
        )}

        {/* Remote Work Option */}
        <Grid size={{ xs: 12 }}>
          <FormControlLabel
            control={
              <Checkbox
                checked={isRemote}
                onChange={handleRemoteToggle}
                disabled={disabled}
                color="primary"
              />
            }
            label="Open to remote work"
          />
        </Grid>

        {/* Location Summary Chips */}
        {(selectedCountry || selectedState || selectedCity || isRemote) && (
          <Grid size={{ xs: 12 }}>
            <Box sx={{ display: 'flex', flexWrap: 'wrap', gap: 1, mt: 1 }}>
              {selectedCountry && (
                <Chip
                  icon={<Public />}
                  label={selectedCountry.name}
                  variant="outlined"
                  color="primary"
                  size="small"
                />
              )}
              {selectedState && (
                <Chip
                  label={selectedState.name}
                  variant="outlined"
                  color="secondary"
                  size="small"
                />
              )}
              {selectedCity && showCity && (
                <Chip
                  icon={<Home />}
                  label={selectedCity.name}
                  variant="outlined"
                  color="default"
                  size="small"
                />
              )}
              {isRemote && (
                <Chip
                  label="Remote"
                  variant="filled"
                  color="success"
                  size="small"
                />
              )}
            </Box>
          </Grid>
        )}
      </Grid>
    </Box>
  );
};

// Demo component to show usage with real data
const LocationInputDemo: React.FC = () => {
  const [location, setLocation] = useState<Partial<Location>>({});
  const [showAdvanced, setShowAdvanced] = useState(false);

  const handleLocationChange = (newLocation: Partial<Location>) => {
    setLocation(newLocation);
    console.log('Location updated:', newLocation);
  };

  // Show some stats about the data
  const totalCountries = Country.getAllCountries().length;
  const usStates = State.getStatesOfCountry('US').length;
  const canadaProvinces = State.getStatesOfCountry('CA').length;

  return (
    <Box sx={{ p: 3, maxWidth: 800, mx: 'auto' }}>
      <Typography variant="h4" gutterBottom align="center" color="primary">
        Location Input with Real Data
      </Typography>
      
      <Typography variant="body2" color="text.secondary" align="center" sx={{ mb: 3 }}>
        Using country-state-city library with {totalCountries} countries, 
        {usStates} US states, {canadaProvinces} Canadian provinces, and thousands of cities
      </Typography>
      
      <Grid container spacing={4}>
        <Grid size={{ xs: 12 }}>
          <Typography variant="h5" gutterBottom>
            Basic Location Input
          </Typography>
          <LocationInput
            value={location}
            onChange={handleLocationChange}
            required
          />
        </Grid>

        <Grid size={{ xs: 12 }}>
          <FormControlLabel
            control={
              <Checkbox
                checked={showAdvanced}
                onChange={(e) => setShowAdvanced(e.target.checked)}
                color="primary"
              />
            }
            label="Show city field"
          />
        </Grid>

        {showAdvanced && (
          <Grid size={{ xs: 12 }}>
            <Typography variant="h5" gutterBottom>
              Advanced Location Input (with City)
            </Typography>
            <LocationInput
              value={location}
              onChange={handleLocationChange}
              showCity
              helperText="Include your city for more specific job matches"
            />
          </Grid>
        )}

        <Grid size={{ xs: 12 }}>
          <Typography variant="h6" gutterBottom>
            Current Location Data:
          </Typography>
          <Box component="pre" sx={{ 
            bgcolor: 'grey.100', 
            p: 2, 
            borderRadius: 1, 
            overflow: 'auto',
            fontSize: '0.875rem'
          }}>
            {JSON.stringify(location, null, 2)}
          </Box>
        </Grid>

        <Grid size={{ xs: 12 }}>
          <Typography variant="body2" color="text.secondary">
            💡 This component uses the country-state-city library which is regularly updated 
            and includes ISO codes, flags, and comprehensive location data.
          </Typography>
        </Grid>
      </Grid>
    </Box>
  );
};

export { LocationInput };