package com.developx.szpitale.load;

import com.developx.szpitale.model.*;
import org.neo4j.driver.Driver;
import org.neo4j.driver.Session;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

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

public class GraphLoader {

    private static final Logger log = LoggerFactory.getLogger(GraphLoader.class);
    private static final int BATCH_SIZE = 1000;

    private final Driver driver;

    public GraphLoader(Driver driver) {
        this.driver = driver;
    }

    public LoadStats load(List<CanonicalDataset> datasets) {
        LoadStats stats = new LoadStats();
        try (Session session = driver.session()) {
            for (CanonicalDataset dataset : datasets) {
                stats.addVoivodeship(dataset.voivodeship());
                loadHospitals(session, dataset.hospitals(), stats);
                loadPeople(session, dataset.people(), stats);
                loadRoles(session, dataset.roles(), stats);
                loadAffiliations(session, dataset.affiliations(), stats);
                loadMandates(session, dataset.mandates(), stats);
                loadOrganizations(session, dataset.organizations(), stats);
                loadCompanies(session, dataset.companies(), stats);
                loadPersonOrganizationLinks(session, dataset.personOrganizationLinks(), stats);
                loadOrganizationLinks(session, dataset.organizationLinks(), stats);
                loadPersonCompanyLinks(session, dataset.personCompanyLinks(), stats);
            }
        }
        return stats;
    }

    private void loadHospitals(Session session, List<Hospital> hospitals, LoadStats stats) {
        List<Map<String, Object>> batch = new ArrayList<>();
        for (Hospital h : hospitals) {
            batch.add(Map.of(
                    "id", h.id(),
                    "name", h.name(),
                    "shortName", h.shortName(),
                    "city", h.city(),
                    "voivodeship", h.voivodeship(),
                    "legalForm", h.legalForm() != null ? h.legalForm().name() : null,
                    "foundingBody", h.foundingBody(),
                    "supervisoryBodyType", h.supervisoryBodyType() != null ? h.supervisoryBodyType().name() : null,
                    "website", h.website(),
                    "sourceUrls", h.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS h " +
                    "MERGE (hospital:Hospital {id: h.id}) " +
                    "SET hospital += {name:h.name, shortName:h.shortName, city:h.city, " +
                    "    voivodeship:h.voivodeship, legalForm:h.legalForm, " +
                    "    foundingBody:h.foundingBody, supervisoryBodyType:h.supervisoryBodyType, " +
                    "    website:h.website} " +
                    "MERGE (v:Voivodeship {slug: h.voivodeship}) " +
                    "MERGE (hospital)-[:W_WOJEWODZTWIE]->(v)";
            session.run(cypher, Map.of("batch", subList));
            stats.addHospitals(subList.size());
        }
    }

    private void loadPeople(Session session, List<Person> people, LoadStats stats) {
        List<Map<String, Object>> batch = new ArrayList<>();
        for (Person p : people) {
            batch.add(Map.of(
                    "id", p.id(),
                    "fullName", p.fullName(),
                    "displayName", p.displayName(),
                    "titles", p.titles() != null ? p.titles() : Collections.emptyList(),
                    "sourceUrls", p.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS p " +
                    "MERGE (person:Person {id: p.id}) " +
                    "SET person += {fullName:p.fullName, displayName:p.displayName, " +
                    "    titles:p.titles, sourceUrls:p.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addPeople(subList.size());
        }
    }

    private void loadRoles(Session session, List<Role> roles, LoadStats stats) {
        Map<String, List<Role>> byRelType = new LinkedHashMap<>();
        for (Role r : roles) {
            String relType = roleTypeToRelName(r);
            byRelType.computeIfAbsent(relType, k -> new ArrayList<>()).add(r);
        }

        for (Map.Entry<String, List<Role>> entry : byRelType.entrySet()) {
            String relType = entry.getKey();
            List<Role> items = entry.getValue();

            for (int i = 0; i < items.size(); i += BATCH_SIZE) {
                List<Role> subList = items.subList(i, Math.min(i + BATCH_SIZE, items.size()));
                String cypher = "UNWIND $batch AS r " +
                        "MATCH (h:Hospital {id: r.hospitalId}) " +
                        "MATCH (p:Person {id: r.personId}) " +
                        "MERGE (h)-[rrel:`" + escapeBacktick(relType) + "`]->(p) " +
                        "SET rrel += {roleType: r.roleType, roleLabel: r.roleLabel, " +
                        "    organ: r.organ, status: r.status, " +
                        "    validFrom: r.validFrom, validTo: r.validTo, note: r.note}";
                session.run(cypher, Map.of("batch", subList.stream()
                        .map(r -> Map.of(
                                "hospitalId", r.hospitalId(),
                                "personId", r.personId(),
                                "roleType", r.roleType().name(),
                                "roleLabel", r.roleLabel(),
                                "organ", r.organ().name(),
                                "status", r.status().name(),
                                "validFrom", r.validFrom(),
                                "validTo", r.validTo(),
                                "note", r.note()
                        ))
                        .collect(Collectors.toList())));
                stats.addRoles(subList.size());
            }
        }
    }

    private void loadAffiliations(Session session, List<Affiliation> affiliations, LoadStats stats) {
        for (int i = 0; i < affiliations.size(); i += BATCH_SIZE) {
            List<Affiliation> subList = affiliations.subList(
                    i, Math.min(i + BATCH_SIZE, affiliations.size()));
            String cypher = "UNWIND $batch AS a " +
                    "MATCH (p:Person {id: a.personId}) " +
                    "MERGE (party:Party {name: a.name}) SET party.type = a.type " +
                    "MERGE (p)-[m:CZLONEK_PARTII]->(party) " +
                    "SET m += {confidence: a.confidence, note: a.note, sourceUrls: a.sourceUrls}";
            session.run(cypher, Map.of("batch", subList.stream()
                    .map(a -> Map.of(
                            "personId", a.personId(),
                            "name", a.name(),
                            "type", a.type().name(),
                            "confidence", a.confidence().name(),
                            "note", a.note(),
                            "sourceUrls", a.sourceUrls()
                    ))
                    .collect(Collectors.toList())));
            stats.addAffiliations(subList.size());
        }
    }

    private void loadMandates(Session session, List<Mandate> mandates, LoadStats stats) {
        if (mandates.isEmpty()) return;

        for (int i = 0; i < mandates.size(); i += BATCH_SIZE) {
            List<Mandate> subList = mandates.subList(
                    i, Math.min(i + BATCH_SIZE, mandates.size()));
            String cypher = "UNWIND $batch AS m " +
                    "MATCH (p:Person {id: m.personId}) " +
                    "MERGE (g:GovBody {name: m.body}) " +
                    "    SET g.type = m.mandateType, g.voivodeship = m.voivodeship, g.term = m.term " +
                    "MERGE (p)-[:PELNI_MANDAT {mandateType: m.mandateType, term: m.term}]->(g)";
            session.run(cypher, Map.of("batch", subList.stream()
                    .map(m -> Map.of(
                            "personId", m.personId(),
                            "body", m.body(),
                            "mandateType", m.mandateType().name(),
                            "term", m.term(),
                            "voivodeship", m.voivodeship(),
                            "sourceUrls", m.sourceUrls()
                    ))
                    .collect(Collectors.toList())));
            stats.addMandates(subList.size());
        }
    }

    private void loadOrganizations(Session session, List<Organization> organizations, LoadStats stats) {
        if (organizations.isEmpty()) return;

        List<Map<String, Object>> batch = new ArrayList<>();
        for (Organization o : organizations) {
            batch.add(Map.of(
                    "id", o.id(),
                    "name", o.name(),
                    "krs", o.krs(),
                    "nip", o.nip(),
                    "sourceUrls", o.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS o " +
                    "MERGE (org:Organization {id: o.id}) " +
                    "SET org += {name: o.name, krs: o.krs, nip: o.nip, sourceUrls: o.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addOrganizations(subList.size());
        }
    }

    private void loadCompanies(Session session, List<Company> companies, LoadStats stats) {
        if (companies.isEmpty()) return;

        List<Map<String, Object>> batch = new ArrayList<>();
        for (Company c : companies) {
            batch.add(Map.of(
                    "id", c.id(),
                    "name", c.name(),
                    "krs", c.krs(),
                    "nip", c.nip(),
                    "ceidgId", c.ceidgId(),
                    "sourceUrls", c.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS c " +
                    "MERGE (comp:Company {id: c.id}) " +
                    "SET comp += {name: c.name, krs: c.krs, nip: c.nip, ceidgId: c.ceidgId, sourceUrls: c.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addCompanies(subList.size());
        }
    }

    private void loadPersonOrganizationLinks(Session session, List<PersonOrganizationLink> links, LoadStats stats) {
        if (links.isEmpty()) return;

        List<Map<String, Object>> batch = new ArrayList<>();
        for (PersonOrganizationLink l : links) {
            batch.add(Map.of(
                    "personId", l.personId(),
                    "organizationId", l.organizationId(),
                    "roleLabel", l.roleLabel(),
                    "basis", l.basis().name(),
                    "sourceUrls", l.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS l " +
                    "MATCH (p:Person {id: l.personId}), (org:Organization {id: l.organizationId}) " +
                    "MERGE (p)-[m:CZLONEK_ORGANIZACJI]->(org) " +
                    "SET m += {roleLabel: l.roleLabel, basis: l.basis, sourceUrls: l.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addPersonOrganizationLinks(subList.size());
        }
    }

    private void loadOrganizationLinks(Session session, List<OrganizationLink> links, LoadStats stats) {
        if (links.isEmpty()) return;

        List<Map<String, Object>> batch = new ArrayList<>();
        for (OrganizationLink l : links) {
            batch.add(Map.of(
                    "fromId", l.fromOrganizationId(),
                    "toId", l.toOrganizationId(),
                    "basis", l.basis().name(),
                    "note", l.note(),
                    "sourceUrls", l.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS l " +
                    "MATCH (a:Organization {id: l.fromId}), (b:Organization {id: l.toId}) " +
                    "MERGE (a)-[m:POWIAZANA_Z]->(b) " +
                    "SET m += {basis: l.basis, note: l.note, sourceUrls: l.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addOrganizationLinks(subList.size());
        }
    }

    private void loadPersonCompanyLinks(Session session, List<PersonCompanyLink> links, LoadStats stats) {
        if (links.isEmpty()) return;

        List<Map<String, Object>> batch = new ArrayList<>();
        for (PersonCompanyLink l : links) {
            batch.add(Map.of(
                    "personId", l.personId(),
                    "companyId", l.companyId(),
                    "roleLabel", l.roleLabel(),
                    "basis", l.basis().name(),
                    "sourceUrls", l.sourceUrls()
            ));
        }

        for (int i = 0; i < batch.size(); i += BATCH_SIZE) {
            List<Map<String, Object>> subList = batch.subList(i, Math.min(i + BATCH_SIZE, batch.size()));
            String cypher = "UNWIND $batch AS l " +
                    "MATCH (p:Person {id: l.personId}), (comp:Company {id: l.companyId}) " +
                    "MERGE (p)-[m:POWIAZANY_Z_FIRMA]->(comp) " +
                    "SET m += {roleLabel: l.roleLabel, basis: l.basis, sourceUrls: l.sourceUrls}";
            session.run(cypher, Map.of("batch", subList));
            stats.addPersonCompanyLinks(subList.size());
        }
    }

    private String roleTypeToRelName(Role r) {
        String organ = r.organ().name();
        String role = r.roleType().name();
        
        if ("DYREKCJA".equals(organ) && "DYREKTOR".equals(role)) return "DYREKTOR";
        if ("DYREKCJA".equals(organ) && ("ZASTEPCA_DYREKTORA".equals(role) || "GLOWNY_KSIEGOWY".equals(role) || "PRZELOZONA_PIELEGNIARKOW".equals(role))) return "ZASTEPCA_DYREKTORA";
        if ("PRZEWODNICZACY_ORGANU".equals(role)) return "PRZEWODNICZY_ORGANOWI";
        return "CZLONEK_ORGANU";
    }

    private String escapeBacktick(String s) {
        return s.replace("`", "\\`");
    }

    public static class LoadStats {
        public int hospitalCount;
        public int personCount;
        public int roleCount;
        public int affiliationCount;
        public int mandateCount;
        public int organizationCount;
        public int companyCount;
        public int personOrganizationLinkCount;
        public int organizationLinkCount;
        public int personCompanyLinkCount;
        public final List<String> voivodeships = new ArrayList<>();

        public void addVoivodeship(String v) { voivodeships.add(v); }
        public void addHospitals(int count) { hospitalCount += count; }
        public void addPeople(int count) { personCount += count; }
        public void addRoles(int count) { roleCount += count; }
        public void addAffiliations(int count) { affiliationCount += count; }
        public void addMandates(int count) { mandateCount += count; }
        public void addOrganizations(int count) { organizationCount += count; }
        public void addCompanies(int count) { companyCount += count; }
        public void addPersonOrganizationLinks(int count) { personOrganizationLinkCount += count; }
        public void addOrganizationLinks(int count) { organizationLinkCount += count; }
        public void addPersonCompanyLinks(int count) { personCompanyLinkCount += count; }

        @Override
        public String toString() {
            return String.format(
                    "Load stats - voivodeships: %s, hospitals: %d, people: %d, roles: %d, affiliations: %d, mandates: %d, organizations: %d, companies: %d, personOrgLinks: %d, orgLinks: %d, personCompanyLinks: %d",
                    String.join(", ", voivodeships), hospitalCount, personCount, roleCount, affiliationCount, mandateCount,
                    organizationCount, companyCount, personOrganizationLinkCount, organizationLinkCount, personCompanyLinkCount
            );
        }
    }
}