PKHeX/PKHeX.Core/Legality/RNG/Locks/TeamLockResult.cs

using System.Collections.Generic;

namespace PKHeX.Core;

/// <summary>
/// Analyzes a <see cref="TeamLock"/> to determine if the provided parameters are a valid end product.
/// </summary>
/// <remarks>
/// When generating a Trainer Team for <see cref="GameVersion.CXD"/>, the game must generate PID values that match the programmed team specifications.
/// The game is 'locked' into a PID generating loop until a valid PID is created, after which the routine is unlocked and proceeds to generate the next member.
/// These locks cause the <see cref="PKM.PID"/> of the current <see cref="PKM"/> to be rerolled until the requisite lock is satisfied.
/// This locking is the same as Method H/J/K and Gen3 Cute Charm by extension, with the additional restriction of not allowing shinies.
/// <see cref="PKM.Nature"/> locks require a certain <see cref="Nature"/>, which is derived from the <see cref="PKM.PID"/>.
/// <see cref="PKM.Gender"/> locks require a certain gender value, which is derived from the <see cref="PKM.PID"/> and <see cref="PersonalInfo.Gender"/> ratio.
/// <see cref="PKM.IsShiny"/> locks require the member to not be shiny. This is enforced for non-shadow members for both games, and for shadow members in <see cref="GameVersion.XD"/>.
/// </remarks>
public sealed class TeamLockResult
{
    /// <summary>
    /// Seed prior to generating the team per <see cref="Specifications"/>.
    /// </summary>
    public readonly uint OriginSeed;

    /// <summary>
    /// Indicates if the input parameters can be generated by the <see cref="Specifications"/> provided.
    /// </summary>
    public readonly bool Valid;

    /// <summary>
    /// NPC Team data containing a list of <see cref="NPCLock"/> members.
    /// </summary>
    internal readonly TeamLock Specifications;

    /// <summary>
    /// Frame the <see cref="OriginSeed"/> is from, based on frames reversed from the input seed provided.
    /// </summary>
    private int OriginFrame;

    /// <summary>
    /// Required CPU Trainer Shiny Value
    /// </summary>
    /// <remarks>
    /// If this value is >= 0, the CPU Trainer Shiny Value must be equal to this value as a <see cref="SeedFrame"/> skipped over a matching interrupt frame.
    /// If this value is <see cref="NOT_FORCED"/>, the CPU Trainer Shiny Value can be anything (except matching any of the <see cref="SeedFrame"/> result members.
    /// </remarks>
    private uint RCSV = NOT_FORCED;

    /// <summary>
    /// Player Trainer Shiny Value
    /// </summary>
    /// <remarks>Only used by <see cref="GameVersion.XD"/> encounters, which disallow shiny shadow members for both player &amp; CPU TSVs.</remarks>
    private readonly uint TSV;

    private readonly Stack<NPCLock> Locks;
    private readonly FrameCache Cache;
    // only save a list of valid nodes we've visited to save memory while we recursively search for a full team.
    private readonly Stack<SeedFrame> Team;

    private const uint NOT_FORCED = uint.MaxValue;

    internal TeamLockResult(TeamLock teamSpec, uint originSeed, uint tsv)
    {
        Locks = new Stack<NPCLock>((Specifications = teamSpec).Locks);
        Team = new Stack<SeedFrame>(Locks.Count);
        Cache = new FrameCache(XDRNG.Prev2(originSeed));
        TSV = tsv;

        Valid = FindLockSeed();
        if (Valid)
            OriginSeed = Cache.GetSeed(OriginFrame);
    }

    /// <summary>
    /// Depth-first search traversal which finds a possible origin for the <see cref="Specifications"/>.
    /// </summary>
    /// <param name="frame">Frame at which the search starts/continues at.</param>
    /// <param name="prior">Prior <see cref="NPCLock"/> data. If this is the last lock in the CPU Team, this is null.</param>
    /// <returns>True if the <see cref="Specifications"/> are valid.</returns>
    private bool FindLockSeed(int frame = 0, NPCLock? prior = null)
    {
        if (Locks.Count == 0) // full team reverse-generated
            return VerifyNPC(frame);

        var current = Locks.Pop();
        var locks = GetPossibleLocks(frame, current, prior);
        foreach (var l in locks)
        {
            Team.Push(l); // possible match
            if (FindLockSeed(l.FrameID, current))
                return true; // all locks are satisfied
            Team.Pop(); // no match, remove
        }

        Locks.Push(current); // return the lock, lock is impossible
        return false;
    }

    /// <summary>
    /// Generates a list of frames the <see cref="current"/> lock data can be generated at.
    /// </summary>
    /// <param name="ctr">Starting frame for the traversal.</param>
    /// <param name="current">Current lock criteria to satisfy. Used to find valid <see cref="SeedFrame"/> results to yield.</param>
    /// <param name="prior">Prior lock criteria. Used for determining when the traversal stops.</param>
    /// <returns>List of possible locks for the provided input.</returns>
    private IEnumerable<SeedFrame> GetPossibleLocks(int ctr, NPCLock current, NPCLock? prior)
    {
        if (prior?.Shadow != false)
            return GetSingleLock(ctr, current);

        return GetAllLocks(ctr, current, (NPCLock)prior);
    }

    /// <summary>
    /// Returns a single <see cref="SeedFrame"/> as the <see cref="current"/> lock must match precisely.
    /// </summary>
    /// <param name="ctr">Starting frame for the traversal.</param>
    /// <param name="current">Current lock criteria to satisfy. Used to find valid <see cref="SeedFrame"/> results to yield.</param>
    private IEnumerable<SeedFrame> GetSingleLock(int ctr, NPCLock current)
    {
        uint pid = (Cache[ctr + 1] << 16) | Cache[ctr];
        if (current.MatchesLock(pid))
            yield return new SeedFrame(pid, ctr + current.FramesConsumed);
        else
            yield break;

        // Reaching here means the single lock didn't cut it. Maybe the frame before it was an anti-shiny reroll?

        // Track if we ever require the CPU Trainer Shiny Value to be a value for a shiny skip.
        // We need to un-set this flag if future frames don't pan out.
        bool forcedOT = false;

        int start = 2;
        while (true)
        {
            var upper = Cache[start + 1];
            var lower = Cache[start];
            // uint cid = upper << 16 | lower;
            var sv = (upper ^ lower) >> 3;
            if (sv == TSV) // XD shiny checks all opponent PKM, even non-shadow.
            {
                // Anti-shiny rerolled! This is a possible frame.
            }
            else if (RCSV != NOT_FORCED) // CPU shiny value is required for a previous lock
            {
                if (sv == RCSV)
                {
                    // No CPU shiny value forced yet. Lets try to skip this lock by requiring the eventual OT to get this shiny.
                    RCSV = sv;
                    forcedOT = true;
                    continue; // don't break
                }

                if (forcedOT) // current call to this method had forced the OT; clear the forced OT before breaking.
                    RCSV = NOT_FORCED;
                yield break; // Since we can't skip this interrupt, we're done.
            }
            // Yield the final rerolled pid instead of the bad anti-shiny (metadata/validation).
            yield return new SeedFrame(pid, start + current.FramesConsumed);
            start += 2;
        }
    }

    /// <summary>
    /// Generates <see cref="SeedFrame"/> nodes until the traversal is ended by an interrupt frame that matches the <see cref="prior"/>.
    /// </summary>
    /// <param name="ctr">Starting frame for the traversal.</param>
    /// <param name="current">Current lock criteria to satisfy. Used to find valid <see cref="SeedFrame"/> results to yield.</param>
    /// <param name="prior">Prior lock criteria. Used for determining when the traversal stops.</param>
    /// <returns>List of possible locks for the provided input.</returns>
    /// <remarks>
    /// An "interrupt" signals the end of the traversal.
    /// Any <see cref="SeedFrame"/> afterwards (when generated forward from the CPU Trainer) will use the interrupt rather than the previous <see cref="SeedFrame"/> that was found for the <see cref="prior"/> lock.
    /// </remarks>
    private IEnumerable<SeedFrame> GetAllLocks(int ctr, NPCLock current, NPCLock prior)
    {
        // Since the prior(next) lock is generated 7+2*n frames after, the worst case break is 7 frames after the PID.
        // Continue reversing until a sequential generation case is found.
        int start = ctr;

        // Track if we ever require the CPU Trainer Shiny Value to be a value for a shiny skip.
        // We need to un-set this flag if future frames don't pan out.
        bool forcedOT = false;

        while (true)
        {
            int p7 = ctr - 7;
            if (p7 > start)
            {
                // check for interrupting cpu team cases
                var upper = Cache[p7 + 1];
                var lower = Cache[p7];
                uint cid = (upper << 16) | lower;
                var sv = (upper ^ lower) >> 3;
                if (sv == TSV) // XD shiny checks all opponent PKM, even non-shadow.
                {
                    // This interrupt is ignored! The result is shiny.
                }
                else if (prior.MatchesLock(cid)) // lock matched cpu mon
                {
                    if (RCSV != NOT_FORCED) // CPU shiny value is required for a previous lock
                    {
                        if (sv != RCSV)
                        {
                            if (forcedOT) // current call to this method had forced the OT; clear the forced OT before breaking.
                                RCSV = NOT_FORCED;
                            yield break; // Since we can't skip this interrupt, we're done.
                        }
                    }
                    else // No CPU shiny value forced yet. Lets try to skip this lock by requiring the eventual OT to get this shiny.
                    {
                        RCSV = sv;
                        forcedOT = true;
                        // don't break
                    }
                }
            }
            uint pid = (Cache[ctr + 1] << 16) | Cache[ctr];
            if (current.MatchesLock(pid))
                yield return new SeedFrame(pid, ctr + current.FramesConsumed);

            ctr += 2;
        }
    }

    /// <summary>
    /// Checks to see if the generated <see cref="Team"/> is compatible with the CPU Trainer Data.
    /// </summary>
    /// <param name="ctr">Ending frame of the traversal.</param>
    /// <returns>True if the <see cref="Specifications"/> are valid.</returns>
    private bool VerifyNPC(int ctr)
    {
        var TID16 = Cache[ctr + 1];
        var SID16 = Cache[ctr];
        var CPUSV = (TID16 ^ SID16) >> 3;
        if (RCSV != NOT_FORCED && RCSV != CPUSV)
            return false; // required CPU Trainer's shiny value did not match the required value.

        int pos = Team.Count - 1; // stack can't do a for loop :(
        foreach (var member in Team)
        {
            var pid = member.PID;
            var psv = ((pid & 0xFFFF) ^ (pid >> 16)) >> 3;

            // check for shiny for Trainer -- XD only
            // if (psv == TSV) // XD shiny checks all opponent PKM, even non-shadow.
            //    return false; // no shiny shadow mons
            // we already checked this when re-generating the team

            // check for shiny for CPU
            if (psv == CPUSV)
            {
                if (!Specifications.Locks[pos].Shadow)
                    return false; // no shiny CPU mons
                if (TSV != NOT_FORCED) // XD
                    return false; // no shiny shadow mons
            }
            pos--;
        }

        OriginFrame = ctr + 2;
        return true;
    }
}