PersonCollection

import java.util.*;
import java.util.stream.Collectors;


public class PersonCollectionImpl implements PersonCollection {
    static class PeopleEmailSorted {
        TreeMap<String, Person> people = new TreeMap<>();

        public void add(Person person){
            people.put(person.getEmail(), person);
        }

        public void remove(Person person){
            people.remove(person.getEmail());
        }

        public List<Person> getSorted(){
            List<Person> sorted = new ArrayList<>();
            for (Map.Entry<String, Person> entry : people.entrySet()) {
                sorted.add(entry.getValue());
            }
            return sorted;
        }
    }

    static class NameTownKey {
        final String name;
        final String town;

        public NameTownKey(String name, String town) {
            this.name = name;
            this.town = town;
        }

        @Override
        public int hashCode() {
            return name.hashCode() + 73 * town.hashCode();
        }

        @Override
        public boolean equals(Object obj){
            if(!(obj instanceof NameTownKey)){
                return false;
            }

            NameTownKey other = (NameTownKey) obj;
            return this.name.equals(other.name) && this.town.equals(other.town);
        }
    }

    private Map<String, Person> peopleByEmail;

    //indices
    private Map<String, PeopleEmailSorted> peopleByEmailDomain;
    private Map<NameTownKey, PeopleEmailSorted> peopleByNameAndTown;
    private Map<String, TreeMap<Integer, PeopleEmailSorted>> peopleByAgeAndTown;

    public PersonCollectionImpl(){
        this.peopleByEmail = new HashMap<>();
        this.peopleByEmailDomain = new HashMap<>();
        this.peopleByNameAndTown = new HashMap<>();
        this.peopleByAgeAndTown = new HashMap<>();
    }

    @Override
    public boolean add(String email, String name, int age, String town) {
        if(this.peopleByEmail.containsKey(email)){
            return false;
        }

        Person person = new Person(email, name, age, town);
        this.peopleByEmail.put(email, person);

        addToIndices(person);

        return true;
    }

    @Override
    public int getCount() {
        return this.peopleByEmail.size();
    }

    @Override
    public boolean delete(String email) {
        Person deletedPerson = this.peopleByEmail.remove(email);

        if(deletedPerson != null){
            removeFromIndices(deletedPerson);
            return true;
        }

        return false;
    }

    @Override
    public Person find(String email) {
        return peopleByEmail.get(email);
    }


    @Override
    public Iterable<Person> findAll(String emailDomain) {
        if(emailDomain == null){
            return new ArrayList<>();
        }

        PeopleEmailSorted peopleEmailSorted =  this.peopleByEmailDomain.get(emailDomain);
        if(peopleEmailSorted == null){
            return new ArrayList<>();
        }

        return peopleEmailSorted.getSorted();
    }

    @Override
    public Iterable<Person> findAll(String name, String town) {
        if(name == null || town == null){
            return new ArrayList<>();
        }

        PeopleEmailSorted peopleEmailSorted = this.peopleByNameAndTown.get(new NameTownKey(name, town));
        if(peopleEmailSorted == null){
            return new ArrayList<>();
        }

        return peopleEmailSorted.getSorted();
    }

    @Override
    public Iterable<Person> findAll(int startAge, int endAge) {
        return this.peopleByEmail.values().stream()
                .filter(person -> person.getAge() >= startAge && person.getAge() <= endAge)
                .sorted(Comparator.comparingInt(Person::getAge).thenComparing(Person::getEmail))
                .collect(Collectors.toList());
    }

    @Override
    public Iterable<Person> findAll(int startAge, int endAge, String town) {
//        if(town == null){
//            return new ArrayList<>();
//        }
//
//        List<Person> people = new ArrayList<>();
//        TreeMap<Integer, PeopleEmailSorted> peopleFromTown = peopleByAgeAndTown.get(town);
//        for (Map.Entry<Integer, PeopleEmailSorted> entry : peopleFromTown.entrySet()) {
//            if(entry.getKey() >= startAge && entry.getKey() <= endAge){
//                people.addAll(entry.getValue().getSorted());
//            }
//        }
//        return people;
        return this.peopleByEmail.values().stream()
                .filter(person -> (person.getAge() >= startAge && person.getAge() <= endAge) && person.getTown().equals(town))
                .sorted(Comparator.comparingInt(Person::getAge).thenComparing(Person::getEmail))
                .collect(Collectors.toList());
    }

    private void addToIndices(Person person) {
        //By Email Domain
        String domain = getDomain(person);
        PeopleEmailSorted peopleEmailSorted = peopleByEmailDomain.get(domain);
        if(peopleEmailSorted == null){
            peopleEmailSorted = new PeopleEmailSorted();
            peopleByEmailDomain.put(domain, peopleEmailSorted);
        }
        peopleEmailSorted.add(person);

        //By Name And Town
        NameTownKey nameTownKey = new NameTownKey(person.getName(), person.getTown());
        PeopleEmailSorted peopleByNameAndTownSorted = peopleByNameAndTown.get(nameTownKey);
        if(peopleByNameAndTownSorted == null){
            peopleByNameAndTownSorted = new PeopleEmailSorted();
            peopleByNameAndTown.put(nameTownKey, peopleByNameAndTownSorted);
        }
        peopleByNameAndTownSorted.add(person);

        //By Age And Town
//        TreeMap<Integer, PeopleEmailSorted> peopleByAgeAndTownSorted = peopleByAgeAndTown.get(person.getTown());
//        if(peopleByAgeAndTownSorted == null){
//            peopleByAgeAndTownSorted = new TreeMap<>();
//            peopleByAgeAndTownSorted.put(person.getAge(), new PeopleEmailSorted());
//            peopleByAgeAndTown.put(person.getTown(), peopleByAgeAndTownSorted);
//        }
//        peopleByAgeAndTownSorted.get(person.getAge()).add(person);
        // TODO: add to other indices
    }

    private void removeFromIndices(Person person) {
        //By Email Domain
        String domain = getDomain(person);
        PeopleEmailSorted peopleEmailSorted = peopleByEmailDomain.get(domain);
        if(peopleEmailSorted != null){
            peopleEmailSorted.remove(person);
        }

        //By Name And Town
        NameTownKey nameTownKey = new NameTownKey(person.getName(), person.getTown());
        PeopleEmailSorted peopleByNameAndTownSorted = peopleByNameAndTown.get(nameTownKey);
        if(peopleByNameAndTownSorted != null){
            peopleByNameAndTownSorted.remove(person);
        }

//        //By Age And Town
//        TreeMap<Integer, PeopleEmailSorted> peopleByAgeAndTownSorted = peopleByAgeAndTown.get(person.getTown());
//        if(peopleByAgeAndTownSorted != null){
//            peopleByAgeAndTownSorted.get(person.getAge()).remove(person);
//        }

        // TODO: remove from other indices

    }

    private String getDomain(Person person){
        return person.getEmail().substring(person.getEmail().lastIndexOf('@') + 1);
    }
}