package com.tiedup.remake.prison;

import com.tiedup.remake.core.TiedUpMod;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import net.minecraft.nbt.CompoundTag;
import net.minecraft.nbt.ListTag;
import net.minecraft.nbt.Tag;
import net.minecraft.server.MinecraftServer;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.server.level.ServerPlayer;
import net.minecraft.world.level.saveddata.SavedData;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/**
 * Central manager for all prisoner data.
 *
 * Replaces:
 * - CaptivityStateManager (state tracking)
 * - Prisoner tracking in CellRegistry
 * - Prisoner tracking in CampOwnership
 *
 * This is the SINGLE SOURCE OF TRUTH for prisoner state.
 * All transitions must go through this class.
 */
public class PrisonerManager extends SavedData {

    private static final String DATA_NAME = "tiedup_prisoner_manager";

    // ==================== PRIMARY DATA ====================

    /** Player UUID -> PrisonerRecord */
    private final Map<UUID, PrisonerRecord> prisoners =
        new ConcurrentHashMap<>();

    /** Player UUID -> LaborRecord */
    private final Map<UUID, LaborRecord> laborRecords =
        new ConcurrentHashMap<>();

    /** Player UUID -> RansomRecord */
    private final Map<UUID, RansomRecord> ransomRecords =
        new ConcurrentHashMap<>();

    // ==================== INDEXES ====================

    /** Camp UUID -> Set of prisoner UUIDs */
    private final Map<UUID, Set<UUID>> prisonersByCamp =
        new ConcurrentHashMap<>();

    /** Cell UUID -> Set of prisoner UUIDs */
    private final Map<UUID, Set<UUID>> prisonersByCell =
        new ConcurrentHashMap<>();

    // ==================== STATIC ACCESS ====================

    public static PrisonerManager get(ServerLevel level) {
        return level
            .getDataStorage()
            .computeIfAbsent(
                PrisonerManager::load,
                PrisonerManager::new,
                DATA_NAME
            );
    }

    public static PrisonerManager get(MinecraftServer server) {
        return get(server.overworld());
    }

    @Nullable
    public static PrisonerManager get(ServerPlayer player) {
        if (player.getServer() == null) return null;
        return get(player.getServer());
    }

    // ==================== RECORD ACCESS ====================

    /**
     * Mark this SavedData as dirty so it gets persisted on next save.
     * Exposed for external callers that modify records directly (e.g., clearing stale guard refs).
     */
    public void markDirty() {
        setDirty();
    }

    /**
     * Get or create the prisoner record for a player.
     * Never returns null - creates a FREE record if none exists.
     */
    public PrisonerRecord getRecord(UUID playerId) {
        return prisoners.computeIfAbsent(playerId, id -> new PrisonerRecord());
    }

    /**
     * Get the prisoner record if the player is a prisoner (non-FREE state).
     * Returns null if player is FREE or has no record.
     */
    @Nullable
    public PrisonerRecord getPrisoner(UUID playerId) {
        PrisonerRecord record = prisoners.get(playerId);
        if (record == null || record.getState() == PrisonerState.FREE) {
            return null;
        }
        return record;
    }

    /**
     * Get or create the labor record for a player.
     */
    public LaborRecord getLaborRecord(UUID playerId) {
        return laborRecords.computeIfAbsent(playerId, id -> new LaborRecord());
    }

    /**
     * Get labor record only if one exists. Does not create ghost entries.
     */
    @Nullable
    public LaborRecord getLaborRecordIfExists(UUID playerId) {
        return laborRecords.get(playerId);
    }

    /**
     * Set the labor record for a player.
     */
    public void setLaborRecord(UUID playerId, LaborRecord record) {
        laborRecords.put(playerId, record);
        setDirty();
    }

    /**
     * Get the ransom record for a player (may be null).
     */
    @Nullable
    public RansomRecord getRansomRecord(UUID playerId) {
        return ransomRecords.get(playerId);
    }

    /**
     * Set the ransom record for a player.
     */
    public void setRansomRecord(UUID playerId, @Nullable RansomRecord record) {
        if (record == null) {
            ransomRecords.remove(playerId);
        } else {
            ransomRecords.put(playerId, record);
        }
        setDirty();
    }

    /**
     * Create or get ransom record for a player.
     */
    public RansomRecord getOrCreateRansomRecord(UUID playerId) {
        return ransomRecords.computeIfAbsent(playerId, id ->
            new RansomRecord()
        );
    }

    /**
     * Check if a player has any prisoner data.
     */
    public boolean hasRecord(UUID playerId) {
        PrisonerRecord record = prisoners.get(playerId);
        return record != null && record.getState() != PrisonerState.FREE;
    }

    // ==================== STATE QUERIES ====================

    /**
     * Get the current state for a player.
     */
    public PrisonerState getState(UUID playerId) {
        PrisonerRecord record = prisoners.get(playerId);
        return record != null ? record.getState() : PrisonerState.FREE;
    }

    /**
     * Check if player is captive (captured or imprisoned).
     */
    public boolean isCaptive(UUID playerId) {
        return getState(playerId).isCaptive();
    }

    /**
     * Check if player is imprisoned (in cell or working).
     */
    public boolean isImprisoned(UUID playerId) {
        return getState(playerId).isImprisoned();
    }

    /**
     * Check if player can be targeted by kidnappers.
     */
    public boolean isTargetable(UUID playerId, long currentTime) {
        PrisonerRecord record = prisoners.get(playerId);
        if (record == null) return true;
        return record.isTargetable(currentTime);
    }

    /**
     * Check if player is protected from capture.
     */
    public boolean isProtected(UUID playerId, long currentTime) {
        PrisonerRecord record = prisoners.get(playerId);
        if (record == null) return false;
        return record.isProtected(currentTime);
    }

    // ==================== INDEX QUERIES ====================

    /**
     * Get all prisoners in a camp.
     */
    public Set<UUID> getPrisonersInCamp(UUID campId) {
        Set<UUID> prisoners = prisonersByCamp.get(campId);
        return prisoners != null
            ? new HashSet<>(prisoners)
            : Collections.emptySet();
    }

    /**
     * Get all prisoners in a cell.
     */
    public Set<UUID> getPrisonersInCell(UUID cellId) {
        Set<UUID> prisoners = prisonersByCell.get(cellId);
        return prisoners != null
            ? new HashSet<>(prisoners)
            : Collections.emptySet();
    }

    /**
     * Get prisoner count in a camp.
     */
    public int getPrisonerCountInCamp(UUID campId) {
        Set<UUID> prisoners = prisonersByCamp.get(campId);
        return prisoners != null ? prisoners.size() : 0;
    }

    /**
     * Get prisoner count in a cell.
     */
    public int getPrisonerCountInCell(UUID cellId) {
        Set<UUID> prisoners = prisonersByCell.get(cellId);
        return prisoners != null ? prisoners.size() : 0;
    }

    /**
     * Get all prisoner IDs (for iteration).
     */
    public Set<UUID> getAllPrisonerIds() {
        return prisoners
            .keySet()
            .stream()
            .filter(id -> prisoners.get(id).getState() != PrisonerState.FREE)
            .collect(Collectors.toSet());
    }

    /**
     * Get all prisoners in a specific state.
     */
    public Set<UUID> getPrisonersInState(PrisonerState state) {
        return prisoners
            .entrySet()
            .stream()
            .filter(e -> e.getValue().getState() == state)
            .map(Map.Entry::getKey)
            .collect(Collectors.toSet());
    }

    /**
     * Get prisoners in camp with specific state.
     */
    public Set<UUID> getPrisonersInCampWithState(
        UUID campId,
        PrisonerState state
    ) {
        Set<UUID> campPrisoners = prisonersByCamp.get(campId);
        if (campPrisoners == null) return Collections.emptySet();

        return campPrisoners
            .stream()
            .filter(id -> {
                PrisonerRecord record = prisoners.get(id);
                return record != null && record.getState() == state;
            })
            .collect(Collectors.toSet());
    }

    // ==================== STATE TRANSITIONS ====================

    /**
     * Capture a free player.
     *
     * @param playerId Player to capture
     * @param captorId UUID of the kidnapper
     * @param currentTime Current game time
     * @return true if capture was successful
     */
    public boolean capture(UUID playerId, UUID captorId, long currentTime) {
        PrisonerRecord record = getRecord(playerId);

        if (
            !PrisonerTransition.capture(record, captorId, currentTime, playerId)
        ) {
            return false;
        }

        setDirty();
        return true;
    }

    /**
     * Imprison a captured player in a camp cell.
     *
     * @param playerId Player to imprison
     * @param campId Camp UUID
     * @param cellId Cell UUID
     * @param currentTime Current game time
     * @return true if imprisonment was successful
     */
    public boolean imprison(
        UUID playerId,
        UUID campId,
        UUID cellId,
        long currentTime
    ) {
        PrisonerRecord record = getRecord(playerId);

        if (
            !PrisonerTransition.imprison(
                record,
                campId,
                cellId,
                currentTime,
                playerId
            )
        ) {
            return false;
        }

        // Update indexes
        addToIndex(prisonersByCamp, campId, playerId);
        addToIndex(prisonersByCell, cellId, playerId);

        // Initialize labor record
        LaborRecord labor = getLaborRecord(playerId);
        labor.setPhase(LaborRecord.WorkPhase.IDLE, currentTime);

        setDirty();
        return true;
    }

    /**
     * Extract prisoner from cell for labor.
     *
     * @param playerId Player to extract
     * @param currentTime Current game time
     * @return true if extraction was successful
     */
    public boolean extract(UUID playerId, long currentTime) {
        PrisonerRecord record = getRecord(playerId);

        if (!PrisonerTransition.extract(record, currentTime, playerId)) {
            return false;
        }

        // Remove from cell index (still in camp though)
        UUID cellId = record.getCellId();
        if (cellId != null) {
            removeFromIndex(prisonersByCell, cellId, playerId);
        }

        setDirty();
        return true;
    }

    /**
     * Return prisoner to cell after labor.
     *
     * @param playerId Player to return
     * @param cellId Cell to return to (may be different from original)
     * @param currentTime Current game time
     * @return true if return was successful
     */
    public boolean returnToCell(UUID playerId, UUID cellId, long currentTime) {
        PrisonerRecord record = getRecord(playerId);

        if (!PrisonerTransition.returnToCell(record, currentTime, playerId)) {
            return false;
        }

        // Update cell assignment and index
        record.setCellId(cellId);
        addToIndex(prisonersByCell, cellId, playerId);

        setDirty();
        return true;
    }

    /**
     * Release a prisoner with grace period.
     *
     * @param playerId Player to release
     * @param currentTime Current game time
     * @param gracePeriodTicks Protection duration (default: 6000 = 5 min)
     * @return true if release was successful
     */
    public boolean release(
        UUID playerId,
        long currentTime,
        long gracePeriodTicks
    ) {
        PrisonerRecord record = getRecord(playerId);
        UUID campId = record.getCampId();
        UUID cellId = record.getCellId();

        if (
            !PrisonerTransition.release(
                record,
                currentTime,
                gracePeriodTicks,
                playerId
            )
        ) {
            return false;
        }

        // Remove from indexes
        if (campId != null) {
            removeFromIndex(prisonersByCamp, campId, playerId);
        }
        if (cellId != null) {
            removeFromIndex(prisonersByCell, cellId, playerId);
        }

        // Clear labor and ransom records
        laborRecords.remove(playerId);
        ransomRecords.remove(playerId);

        setDirty();
        return true;
    }

    /**
     * Release a prisoner with default 5-minute grace period.
     *
     * @param playerId Player to release
     * @param currentTime Current game time
     * @return true if release was successful
     */
    public boolean release(UUID playerId, long currentTime) {
        return release(playerId, currentTime, 6000); // 5 minutes default
    }

    /**
     * Transition prisoner from IMPRISONED to WORKING state.
     * Used when extracting for labor.
     *
     * @param playerId Player to transition
     * @param currentTime Current game time
     * @return true if transition was successful
     */
    public boolean transitionToWorking(UUID playerId, long currentTime) {
        return extract(playerId, currentTime);
    }

    /**
     * Mark a prisoner as escaped (goes directly to FREE).
     *
     * @param playerId Player who escaped
     * @param currentTime Current game time
     * @param reason Reason for escape (for logging)
     * @return true if escape was processed
     */
    public boolean escape(UUID playerId, long currentTime, String reason) {
        PrisonerRecord record = getRecord(playerId);
        UUID campId = record.getCampId();
        UUID cellId = record.getCellId();

        if (!PrisonerTransition.escape(record, currentTime, playerId, reason)) {
            return false;
        }

        // Remove from indexes
        if (campId != null) {
            removeFromIndex(prisonersByCamp, campId, playerId);
        }
        if (cellId != null) {
            removeFromIndex(prisonersByCell, cellId, playerId);
        }

        // Clear labor and ransom records
        laborRecords.remove(playerId);
        ransomRecords.remove(playerId);

        setDirty();
        return true;
    }

    /**
     * Expire protection (PROTECTED -> FREE).
     */
    public boolean expireProtection(UUID playerId, long currentTime) {
        PrisonerRecord record = prisoners.get(playerId);
        if (record == null || record.getState() != PrisonerState.PROTECTED) {
            return false;
        }

        PrisonerTransition.forceTransition(
            record,
            PrisonerState.FREE,
            currentTime
        );
        setDirty();
        return true;
    }

    // ==================== RANSOM MANAGEMENT ====================

    /**
     * Create a ransom for a prisoner.
     *
     * @param playerId Prisoner UUID
     * @param totalDebt Total debt amount
     * @param currentTime Current game time
     */
    public void createRansom(UUID playerId, int totalDebt, long currentTime) {
        RansomRecord ransom = getOrCreateRansomRecord(playerId);
        ransom.setTotalDebt(totalDebt);
        setDirty();
    }

    /**
     * Add payment to a prisoner's ransom.
     *
     * @param playerId Prisoner UUID
     * @param amount Payment amount
     * @param contributorId Player who paid (null for labor)
     * @return true if ransom is now fully paid
     */
    public boolean addRansomPayment(
        UUID playerId,
        int amount,
        @Nullable UUID contributorId
    ) {
        RansomRecord ransom = ransomRecords.get(playerId);
        if (ransom == null) return false;

        boolean paid = ransom.addPayment(amount, contributorId);
        setDirty();
        return paid;
    }

    /**
     * Increase prisoner's debt (punishment).
     */
    public void increaseDebt(UUID playerId, int amount) {
        RansomRecord ransom = ransomRecords.get(playerId);
        if (ransom != null) {
            ransom.increaseDebt(amount);
            setDirty();
        }
    }

    // ==================== LABOR MANAGEMENT ====================

    /**
     * Assign a task to a prisoner.
     */
    public boolean assignTask(
        UUID playerId,
        com.tiedup.remake.labor.LaborTask task,
        UUID maidId,
        long currentTime
    ) {
        LaborRecord labor = getLaborRecord(playerId);
        if (!labor.canAssignTask()) {
            return false;
        }

        labor.assignTask(task, maidId, currentTime);
        setDirty();
        return true;
    }

    /**
     * Get the current work phase for a prisoner.
     */
    public LaborRecord.WorkPhase getWorkPhase(UUID playerId) {
        LaborRecord labor = laborRecords.get(playerId);
        return labor != null ? labor.getPhase() : LaborRecord.WorkPhase.IDLE;
    }

    // ==================== INDEX HELPERS ====================

    private void addToIndex(Map<UUID, Set<UUID>> index, UUID key, UUID value) {
        if (key == null) return;
        index
            .computeIfAbsent(key, k -> ConcurrentHashMap.newKeySet())
            .add(value);
    }

    private void removeFromIndex(
        Map<UUID, Set<UUID>> index,
        UUID key,
        UUID value
    ) {
        if (key == null) return;
        Set<UUID> set = index.get(key);
        if (set != null) {
            set.remove(value);
            if (set.isEmpty()) {
                index.remove(key);
            }
        }
    }

    // ==================== CLEANUP ====================

    /**
     * Clean up stale records for offline players.
     * Called periodically.
     */
    public void cleanupOfflinePlayers(
        MinecraftServer server,
        long currentTime,
        long offlineTimeoutTicks
    ) {
        List<UUID> toCleanup = new ArrayList<>();

        for (Map.Entry<UUID, PrisonerRecord> entry : prisoners.entrySet()) {
            UUID playerId = entry.getKey();
            PrisonerRecord record = entry.getValue();

            // Skip FREE players
            if (record.getState() == PrisonerState.FREE) {
                continue;
            }

            // Check if player is online
            ServerPlayer player = server.getPlayerList().getPlayer(playerId);
            if (player != null) {
                continue; // Online, skip
            }

            // Check timeout
            long timeInState = record.getTimeInState(currentTime);
            if (timeInState > offlineTimeoutTicks) {
                toCleanup.add(playerId);
            }
        }

        for (UUID playerId : toCleanup) {
            TiedUpMod.LOGGER.info(
                "[PrisonerManager] Cleaning up offline prisoner: {}",
                playerId.toString().substring(0, 8)
            );
            escape(playerId, currentTime, "offline timeout");
        }
    }

    /**
     * Expire protection for players whose grace period ended.
     */
    public void tickProtectionExpiry(long currentTime) {
        List<UUID> toExpire = new ArrayList<>();

        for (Map.Entry<UUID, PrisonerRecord> entry : prisoners.entrySet()) {
            PrisonerRecord record = entry.getValue();
            if (
                record.getState() == PrisonerState.PROTECTED &&
                currentTime >= record.getProtectionExpiry()
            ) {
                toExpire.add(entry.getKey());
            }
        }

        for (UUID playerId : toExpire) {
            expireProtection(playerId, currentTime);
        }
    }

    // ==================== PERSISTENCE ====================

    @Override
    public @NotNull CompoundTag save(@NotNull CompoundTag tag) {
        // Save prisoners
        ListTag prisonerList = new ListTag();
        for (Map.Entry<UUID, PrisonerRecord> entry : prisoners.entrySet()) {
            CompoundTag prisonerTag = new CompoundTag();
            prisonerTag.putUUID("id", entry.getKey());
            prisonerTag.put("record", entry.getValue().save());
            prisonerList.add(prisonerTag);
        }
        tag.put("prisoners", prisonerList);

        // Save labor records
        ListTag laborList = new ListTag();
        for (Map.Entry<UUID, LaborRecord> entry : laborRecords.entrySet()) {
            CompoundTag laborTag = new CompoundTag();
            laborTag.putUUID("id", entry.getKey());
            laborTag.put("record", entry.getValue().save());
            laborList.add(laborTag);
        }
        tag.put("laborRecords", laborList);

        // Save ransom records
        ListTag ransomList = new ListTag();
        for (Map.Entry<UUID, RansomRecord> entry : ransomRecords.entrySet()) {
            CompoundTag ransomTag = new CompoundTag();
            ransomTag.putUUID("id", entry.getKey());
            ransomTag.put("record", entry.getValue().save());
            ransomList.add(ransomTag);
        }
        tag.put("ransomRecords", ransomList);

        return tag;
    }

    public static PrisonerManager load(CompoundTag tag) {
        PrisonerManager manager = new PrisonerManager();

        // Load prisoners
        if (tag.contains("prisoners")) {
            ListTag prisonerList = tag.getList("prisoners", Tag.TAG_COMPOUND);
            for (int i = 0; i < prisonerList.size(); i++) {
                CompoundTag prisonerTag = prisonerList.getCompound(i);
                UUID id = prisonerTag.getUUID("id");
                PrisonerRecord record = PrisonerRecord.load(
                    prisonerTag.getCompound("record")
                );
                manager.prisoners.put(id, record);

                // Rebuild indexes
                if (record.getCampId() != null) {
                    manager.addToIndex(
                        manager.prisonersByCamp,
                        record.getCampId(),
                        id
                    );
                }
                if (record.getCellId() != null) {
                    manager.addToIndex(
                        manager.prisonersByCell,
                        record.getCellId(),
                        id
                    );
                }
            }
        }

        // Load labor records
        if (tag.contains("laborRecords")) {
            ListTag laborList = tag.getList("laborRecords", Tag.TAG_COMPOUND);
            for (int i = 0; i < laborList.size(); i++) {
                CompoundTag laborTag = laborList.getCompound(i);
                UUID id = laborTag.getUUID("id");
                LaborRecord record = LaborRecord.load(
                    laborTag.getCompound("record")
                );
                manager.laborRecords.put(id, record);
            }
        }

        // Load ransom records
        if (tag.contains("ransomRecords")) {
            ListTag ransomList = tag.getList("ransomRecords", Tag.TAG_COMPOUND);
            for (int i = 0; i < ransomList.size(); i++) {
                CompoundTag ransomTag = ransomList.getCompound(i);
                UUID id = ransomTag.getUUID("id");
                RansomRecord record = RansomRecord.load(
                    ransomTag.getCompound("record")
                );
                manager.ransomRecords.put(id, record);
            }
        }

        return manager;
    }

    // ==================== DEBUG ====================

    public String toDebugString() {
        StringBuilder sb = new StringBuilder();
        sb.append("PrisonerManager:\n");
        sb.append("  Total records: ").append(prisoners.size()).append("\n");
        sb
            .append("  Active prisoners: ")
            .append(getAllPrisonerIds().size())
            .append("\n");

        for (PrisonerState state : PrisonerState.values()) {
            int count = getPrisonersInState(state).size();
            if (count > 0) {
                sb
                    .append("  - ")
                    .append(state)
                    .append(": ")
                    .append(count)
                    .append("\n");
            }
        }

        return sb.toString();
    }
}