#!/usr/bin/env python3
"""Generate the exact Magic120Cell sticker-permutation table.

The extractor compiles a temporary, headless compatibility copy of the
MIT-licensed roice3/Magic120Cell geometry/state engine. It assigns a unique
identity to each of the 7,560 sticker slots, applies every cell/sticker move,
and records only the changed destination/source pairs.

Copyright (c) 2016 Roice Nelson
Magic120Cell source license: MIT
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import shutil
import struct
import subprocess
import tempfile
import zlib
from pathlib import Path


CELL_COUNT = 120
STICKERS_PER_CELL = 63
TWIST_STICKER_COUNT = 62
STATE_SIZE = CELL_COUNT * STICKERS_PER_CELL
MOVE_COUNT = CELL_COUNT * TWIST_STICKER_COUNT
RAW_MAGIC = b"HTM120P1"


COMPAT_STDAFX = r"""
#pragma once

#include <GL/gl.h>
#include <algorithm>
#include <array>
#include <assert.h>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <vector>

using __int8 = std::int8_t;
using __int16 = std::int16_t;
using __int32 = std::int32_t;
using __int64 = std::int64_t;
using BOOL = int;
using HDC = void*;
using HGLRC = void*;
using HWND = void*;

#ifndef TRUE
#define TRUE 1
#endif
#ifndef FALSE
#define FALSE 0
#endif
"""


COMPAT_LOADER_HEADER = r"""
#pragma once

#include "puzzle/state.h"
#include "puzzle/twist.h"

class CLoader
{
public:
    CLoader();
    void saveToFile(int, const CState&, const std::vector<STwist>&, bool);
    void saveScrambledFile(const CState&, const std::vector<STwist>&);
    bool loadFromFile(int&, CState&, std::vector<STwist>&);
    bool loadScrambledFile(CState&, std::vector<STwist>&);
};
"""


COMPAT_RENDERER_HEADER = r"""
#pragma once

#include "vectorND.h"

class IRenderable
{
public:
    virtual ~IRenderable() = default;
    virtual void render(const CVector3D& lookFrom, bool forPicking) = 0;
};
"""


COMPAT_SUPPORT_CPP = r"""
#include <stdafx.h>
#include "helper.h"
#include "loader.h"

void CColor::generateRandom()
{
    m_r = getRandomDouble(1.0);
    m_g = getRandomDouble(1.0);
    m_b = getRandomDouble(1.0);
    m_a = 1.0;
}

void CColor::setColorHLS(double h, double l, double s, double a)
{
    // Color values are irrelevant to permutation extraction, but state
    // initialization requires a stable in-range palette.
    m_r = std::fmod(h + s, 1.0);
    m_g = l;
    m_b = std::fmod(h + 0.5, 1.0);
    m_a = a;
}

void CColor::lighten()
{
    m_r = std::min(1.0, m_r * 1.1);
    m_g = std::min(1.0, m_g * 1.1);
    m_b = std::min(1.0, m_b * 1.1);
}

CLoader::CLoader() = default;
void CLoader::saveToFile(int, const CState&, const std::vector<STwist>&, bool) {}
void CLoader::saveScrambledFile(const CState&, const std::vector<STwist>&) {}
bool CLoader::loadFromFile(int&, CState&, std::vector<STwist>&) { return false; }
bool CLoader::loadScrambledFile(CState&, std::vector<STwist>&) { return false; }
"""


EXTRACTOR_CPP = r"""
#include <stdafx.h>
#include "magic120Cell.h"

namespace
{
    constexpr int CellCount = 120;
    constexpr int StickersPerCell = 63;
    constexpr int StateSize = CellCount * StickersPerCell;

    class CInspectableMagic120Cell final : public CMagic120Cell
    {
    public:
        void setUniqueIdentity()
        {
            for (int cell = 0; cell < CellCount; ++cell)
            {
                for (int sticker = 0; sticker < StickersPerCell; ++sticker)
                {
                    m_state.setStickerColorIndex(
                        cell,
                        sticker,
                        cell * StickersPerCell + sticker);
                }
            }
            m_state.commitChanges();
            m_twistHistory.clear();
        }

        int getIdentityAt(const int destination) const
        {
            return m_state.getStickerColorIndex(
                destination / StickersPerCell,
                destination % StickersPerCell);
        }
    };

    void writeU16(std::ofstream& output, const std::uint16_t value)
    {
        output.write(reinterpret_cast<const char*>(&value), sizeof(value));
    }

    void writeU32(std::ofstream& output, const std::uint32_t value)
    {
        output.write(reinterpret_cast<const char*>(&value), sizeof(value));
    }
}

int main(const int argc, const char* argv[])
{
    if (argc != 2)
    {
        std::cerr << "usage: extractor <raw-output>" << std::endl;
        return 2;
    }

    std::ofstream output(argv[1], std::ios::binary | std::ios::trunc);
    if (!output)
    {
        std::cerr << "unable to open output" << std::endl;
        return 3;
    }

    output.write("HTM120P1", 8);
    writeU32(output, CellCount);
    writeU32(output, StickersPerCell);
    writeU32(output, CellCount * (StickersPerCell - 1));
    writeU32(output, StateSize);

    auto puzzle = std::make_unique<CInspectableMagic120Cell>();
    std::uint64_t totalPairs = 0;
    std::uint16_t minimumPairs = UINT16_MAX;
    std::uint16_t maximumPairs = 0;
    for (int cell = 0; cell < CellCount; ++cell)
    {
        for (int sticker = 1; sticker < StickersPerCell; ++sticker)
        {
            puzzle->setUniqueIdentity();

            STwist twist;
            twist.m_cell = cell;
            twist.m_sticker = sticker;
            twist.m_leftClick = false;
            twist.m_viewRotation = false;
            twist.m_slicemask = 1;
            puzzle->startRotate(twist);
            puzzle->finishRotate();

            std::vector<std::pair<std::uint16_t, std::uint16_t>> changed;
            changed.reserve(256);
            for (int destination = 0; destination < StateSize; ++destination)
            {
                const int source = puzzle->getIdentityAt(destination);
                if (source != destination)
                {
                    if (source < 0 || source >= StateSize)
                    {
                        std::cerr << "invalid source identity" << std::endl;
                        return 4;
                    }
                    changed.emplace_back(
                        static_cast<std::uint16_t>(destination),
                        static_cast<std::uint16_t>(source));
                }
            }

            if (changed.empty() || changed.size() > UINT16_MAX)
            {
                std::cerr << "invalid changed-slot count" << std::endl;
                return 5;
            }
            const auto changedCount =
                static_cast<std::uint16_t>(changed.size());
            writeU16(output, changedCount);
            for (const auto& pair : changed)
            {
                writeU16(output, pair.first);
                writeU16(output, pair.second);
            }

            totalPairs += changedCount;
            minimumPairs = std::min(minimumPairs, changedCount);
            maximumPairs = std::max(maximumPairs, changedCount);
        }
        if ((cell + 1) % 10 == 0)
        {
            std::cerr << "extracted cells " << (cell + 1)
                      << " / " << CellCount << std::endl;
        }
    }

    output.flush();
    if (!output)
    {
        std::cerr << "write failed" << std::endl;
        return 6;
    }
    std::cerr << "moves=" << CellCount * (StickersPerCell - 1)
              << " pairs=" << totalPairs
              << " min=" << minimumPairs
              << " max=" << maximumPairs << std::endl;
    return 0;
}
"""


SOURCE_FILES = (
    "Magic120Cell/workFiles/magic120Cell.cpp",
    "Magic120Cell/workFiles/puzzle.cpp",
    "Magic120Cell/workFiles/cell.cpp",
    "Magic120Cell/workFiles/cellDodec.cpp",
    "Magic120Cell/workFiles/primitives.cpp",
    "Magic120Cell/workFiles/settings.cpp",
    "geometryLib/vectorND.cpp",
    "geometryLib/polygon.cpp",
    "geometryLib/gl/displayList.cpp",
    "geometryLib/puzzle/state.cpp",
    "geometryLib/puzzle/twist.cpp",
    "geometryLib/puzzle/twistHistory.cpp",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--donor-root",
        type=Path,
        default=Path(
            "/home/dev/src/Workspaces/HyperTwist/mirrors/permissive/"
            "roice3/Magic120Cell"
        ),
    )
    parser.add_argument("--output-inl", required=True, type=Path)
    parser.add_argument("--output-manifest", required=True, type=Path)
    parser.add_argument("--keep-raw", type=Path)
    return parser.parse_args()


def patch_compatibility_tree(root: Path) -> None:
    (root / "Magic120Cell/stdafx.h").write_text(COMPAT_STDAFX, encoding="utf-8")
    (root / "Magic120Cell/workFiles/loader.h").write_text(
        COMPAT_LOADER_HEADER, encoding="utf-8"
    )
    (root / "geometryLib/renderer.h").write_text(
        COMPAT_RENDERER_HEADER, encoding="utf-8"
    )
    matrix_text = (root / "geometryLib/matrix4D.h").read_text(encoding="utf-8-sig")
    matrix_text = matrix_text.replace(
        "m[i][j] operation= rhs.m[i][j];",
        "m[i][j] operation ## = rhs.m[i][j];",
    )
    matrix_text = matrix_text.replace(
        "m[i][j] operation= a;",
        "m[i][j] operation ## = a;",
    )
    matrix_text = matrix_text.replace(
        "DO_OP_MATRIX_INPLACE(+);",
        """for( int i=0; i<4; i++ )
\t\t\tfor( int j=0; j<4; j++ )
\t\t\t\tm[i][j] += rhs.m[i][j];
\t\treturn( *this );""",
    )
    matrix_text = matrix_text.replace(
        "DO_OP_M_SCALAR_INPLACE(*);",
        """for( int i=0; i<4; i++ )
\t\t\tfor( int j=0; j<4; j++ )
\t\t\t\tm[i][j] *= a;
\t\treturn( *this );""",
    )
    (root / "geometryLib/matrix4D.h").write_text(matrix_text, encoding="utf-8")
    (root / "geometryLib/matrix4d.h").write_text(matrix_text, encoding="utf-8")
    (root / "geometryLib/gl/displayList.h").write_text(
        """#pragma once

#include <GL/gl.h>

class CDisplayList
{
public:
    CDisplayList();
    ~CDisplayList();
    CDisplayList(const CDisplayList&);
    CDisplayList& operator=(const CDisplayList&);
    void start(bool execute = true);
    void end() const;
    bool isRecorded() const;
    bool call() const;
    void clear();

private:
    GLuint m_dl;
};
""",
        encoding="utf-8",
    )

    vector_header = root / "geometryLib/vectorND.h"
    vector_text = vector_header.read_text(encoding="utf-8-sig")
    vector_text = vector_text.replace(
        "m_components[i] operation= rhs.m_components[i];",
        "m_components[i] operation ## = rhs.m_components[i];",
    )
    vector_text = vector_text.replace(
        "m_components[i] operation= a;",
        "m_components[i] operation ## = a;",
    )
    old_nan_check = """inline bool isNaN( double value )
{
\tint floatType = _fpclass( value );
\treturn
\t\t_FPCLASS_SNAN == floatType ||
\t\t_FPCLASS_QNAN == floatType ||
\t\t_FPCLASS_NINF == floatType ||
\t\t_FPCLASS_PINF == floatType;
}"""
    if old_nan_check not in vector_text:
        raise RuntimeError("expected donor floating-point classification block missing")
    vector_text = vector_text.replace(
        old_nan_check,
        """inline bool isNaN( double value )
{
\treturn !std::isfinite( value );
}""",
    )
    vector_header.write_text(vector_text, encoding="utf-8")

    polygon_path = root / "geometryLib/polygon.cpp"
    polygon_text = polygon_path.read_text(encoding="cp1252")
    polygon_text = polygon_text.replace(
        "s.m_p2 * finiteScale : s.m_p1",
        "CVector3D( s.m_p2 * finiteScale ) : s.m_p1",
    )
    polygon_text = polygon_text.replace(
        "s.m_p1 * finiteScale : s.m_p2",
        "CVector3D( s.m_p1 * finiteScale ) : s.m_p2",
    )
    polygon_path.write_text(polygon_text, encoding="utf-8")

    replacements = {
        "Magic120Cell/workFiles/cellDodec.cpp": (
            ('#include "CellDodec.h"', '#include "cellDodec.h"'),
            ("__super::operator = ( rhs );", "Cell::operator = ( rhs );"),
            ("__super::getSticker( sticker )", "Cell::getSticker( sticker )"),
        ),
        "Magic120Cell/workFiles/magic120Cell.cpp": (
            ('#include "magic120cell.h"', '#include "magic120Cell.h"'),
            ("__super::render( lookFrom, forPicking );", "CPuzzle::render( lookFrom, forPicking );"),
            ("__super::resetView();", "CPuzzle::resetView();"),
        ),
    }
    for relative_path, path_replacements in replacements.items():
        path = root / relative_path
        text = path.read_text(encoding="utf-8-sig")
        for old, new in path_replacements:
            if old not in text:
                raise RuntimeError(f"expected donor source token missing: {old}")
            text = text.replace(old, new)
        path.write_text(text, encoding="utf-8")


def compile_and_extract(root: Path, raw_output: Path) -> str:
    support_cpp = root / "compat_support.cpp"
    extractor_cpp = root / "extractor.cpp"
    executable = root / "magic120cell-extractor"
    support_cpp.write_text(COMPAT_SUPPORT_CPP, encoding="utf-8")
    extractor_cpp.write_text(EXTRACTOR_CPP, encoding="utf-8")

    command = [
        "clang++",
        "-std=c++14",
        "-O2",
        "-g",
        "-fno-omit-frame-pointer",
        "-fsanitize=address,undefined",
        "-ffunction-sections",
        "-fdata-sections",
        "-Wno-unknown-pragmas",
        "-I",
        str(root / "Magic120Cell"),
        "-I",
        str(root / "Magic120Cell/workFiles"),
        "-I",
        str(root / "geometryLib"),
        str(extractor_cpp),
        str(support_cpp),
        *(str(root / relative_path) for relative_path in SOURCE_FILES),
        "-Wl,--gc-sections",
        "-lGL",
        "-o",
        str(executable),
    ]
    compile_result = subprocess.run(
        command, check=False, capture_output=True, text=True
    )
    if compile_result.returncode != 0:
        raise RuntimeError(
            "Magic120Cell extractor compile failed:\n"
            f"{compile_result.stdout}\n{compile_result.stderr}"
        )

    extract_result = subprocess.run(
        [str(executable), str(raw_output)],
        check=False,
        capture_output=True,
        text=True,
    )
    if extract_result.returncode != 0:
        raise RuntimeError(
            "Magic120Cell extractor execution failed:\n"
            f"exit_code={extract_result.returncode}\n"
            f"{extract_result.stdout}\n{extract_result.stderr}"
        )
    return extract_result.stderr.strip()


def validate_raw_table(raw: bytes) -> dict[str, int]:
    if len(raw) < 24 or raw[:8] != RAW_MAGIC:
        raise RuntimeError("raw permutation table has an invalid header")
    cell_count, sticker_count, move_count, state_size = struct.unpack_from(
        "<IIII", raw, 8
    )
    if (
        cell_count != CELL_COUNT
        or sticker_count != STICKERS_PER_CELL
        or move_count != MOVE_COUNT
        or state_size != STATE_SIZE
    ):
        raise RuntimeError("raw permutation table dimensions are invalid")

    cursor = 24
    minimum_pairs = 1 << 30
    maximum_pairs = 0
    total_pairs = 0
    for _ in range(move_count):
        if cursor + 2 > len(raw):
            raise RuntimeError("raw permutation table ended before a move header")
        pair_count = struct.unpack_from("<H", raw, cursor)[0]
        cursor += 2
        byte_count = pair_count * 4
        if pair_count == 0 or cursor + byte_count > len(raw):
            raise RuntimeError("raw permutation table contains an invalid move")

        destinations: set[int] = set()
        sources: set[int] = set()
        for pair_index in range(pair_count):
            destination, source = struct.unpack_from(
                "<HH", raw, cursor + pair_index * 4
            )
            if (
                destination >= state_size
                or source >= state_size
                or destination in destinations
                or source in sources
            ):
                raise RuntimeError("raw move is not a bijection over changed slots")
            destinations.add(destination)
            sources.add(source)
        if destinations != sources:
            raise RuntimeError("raw move does not close over its changed-slot set")

        cursor += byte_count
        minimum_pairs = min(minimum_pairs, pair_count)
        maximum_pairs = max(maximum_pairs, pair_count)
        total_pairs += pair_count
    if cursor != len(raw):
        raise RuntimeError("raw permutation table contains trailing bytes")
    return {
        "minimum_pairs_per_move": minimum_pairs,
        "maximum_pairs_per_move": maximum_pairs,
        "total_pairs": total_pairs,
    }


def write_inl(path: Path, compressed: bytes, raw_size: int, sha256: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    encoded = base64.b64encode(compressed).decode("ascii")
    wrapped = "\n".join(
        f'        TEXT("{encoded[index:index + 120]}"),'
        for index in range(0, len(encoded), 120)
    )
    content = f"""#pragma once

// Generated from roice3/Magic120Cell (MIT), Copyright (c) 2016 Roice Nelson.
// Raw table SHA-256: {sha256}
namespace HyperTwistMagic120CellGenerated
{{
    constexpr int32 RawPermutationTableSize = {raw_size};
    constexpr int32 CompressedPermutationTableSize = {len(compressed)};
    constexpr const TCHAR* CompressedPermutationTableBase64Chunks[] =
    {{
{wrapped}
    }};
}}
"""
    path.write_text(content, encoding="utf-8", newline="\n")


def main() -> int:
    args = parse_args()
    donor_root = args.donor_root.resolve()
    if not (donor_root / "Magic120Cell/workFiles/magic120Cell.cpp").is_file():
        raise RuntimeError(f"Magic120Cell donor root is incomplete: {donor_root}")

    with tempfile.TemporaryDirectory(prefix="hypertwist-magic120cell-") as temporary:
        compatibility_root = Path(temporary) / "source"
        shutil.copytree(donor_root, compatibility_root)
        patch_compatibility_tree(compatibility_root)
        raw_output = Path(temporary) / "magic120cell-permutations.raw"
        extractor_summary = compile_and_extract(compatibility_root, raw_output)
        raw = raw_output.read_bytes()
        table_stats = validate_raw_table(raw)

        compressed = zlib.compress(raw, level=9)
        raw_sha256 = hashlib.sha256(raw).hexdigest()
        write_inl(args.output_inl, compressed, len(raw), raw_sha256)

        if args.keep_raw is not None:
            args.keep_raw.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(raw_output, args.keep_raw)

    manifest = {
        "schema_version": "hypertwist/magic120cell-permutation-table/v1",
        "source_repository": "roice3/Magic120Cell",
        "source_license": "MIT",
        "source_copyright": "Copyright (c) 2016 Roice Nelson",
        "cell_count": CELL_COUNT,
        "stickers_per_cell": STICKERS_PER_CELL,
        "state_size": STATE_SIZE,
        "move_count": MOVE_COUNT,
        "raw_bytes": len(raw),
        "compressed_bytes": len(compressed),
        "raw_sha256": raw_sha256,
        "extractor_summary": extractor_summary,
        **table_stats,
    }
    args.output_manifest.parent.mkdir(parents=True, exist_ok=True)
    args.output_manifest.write_text(
        json.dumps(manifest, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(manifest, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())