import os
import jijmodeling as jm
import glob
from ommx.artifact import ArtifactBuilder
from model import build_ip_formulation
from sol_reader import parse_solution_zfx
from ommx_quantum_benchmarks.qoblib.definitions import (
    QOBLIB_AUTHORS,
    QOBLIB_AUTHORS_STR,
    LICENSE,
)


def _pick_solution_file(sol_dir: str, base: str) -> str | None:
    """Pick first existing solution file for a basename."""
    candidates = [
        os.path.join(sol_dir, f"{base}.opt.sol"),
        os.path.join(sol_dir, f"{base}.best.sol"),
        os.path.join(sol_dir, f"{base}.bst.sol"),
        os.path.join(sol_dir, f"{base}.sol"),
    ]
    for p in candidates:
        if os.path.exists(p):
            return p
    return None


def batch_process(
    sol_root: str = "../../solutions",
    output_directory: str = "./ommx_output",
):
    """
    Process instances from a QBench JSON file and corresponding solution files,
    convert them into OMMX artifacts, and save them to the output directory.

    Parameters:
        sol_root (str): Path to the root solutions directory.
        output_directory (str): Path to save the generated .ommx files.
    """
    os.makedirs(output_directory, exist_ok=True)

    problem = build_ip_formulation()

    processed_count = 0
    error_count = 0

    # scan all sol_root's .sol file and get the base name.
    sol_files = glob.glob(os.path.join(sol_root, "*.sol"))
    bases = set()
    for path in sol_files:
        name = os.path.basename(path)
        for suffix in [".opt.sol", ".best.sol", ".bst.sol", ".sol"]:
            if name.endswith(suffix):
                bases.add(name[: -len(suffix)])
                break

    print(f"Found {len(bases)} bases: {sorted(bases)}")

    def cut_matrix(t: list[list[int]], n: int) -> list[list[int]]:
        if not (5 <= n <= 24):
            raise ValueError("n must be between 5 and 24.")
        return [row[:n] for row in t[:n]]

    your_t_0based = [
        [
            0,
            24,
            43,
            23,
            21,
            41,
            61,
            21,
            20,
            0,
            0,
            0,
            0,
            0,
            40,
            0,
            0,
            0,
            0,
            0,
            24,
            19,
            23,
            0,
        ],
        [
            0,
            0,
            0,
            0,
            21,
            18,
            39,
            23,
            0,
            0,
            40,
            19,
            64,
            19,
            17,
            64,
            80,
            0,
            0,
            22,
            24,
            18,
            19,
            0,
        ],
        [
            16,
            0,
            0,
            20,
            44,
            0,
            42,
            22,
            20,
            0,
            0,
            23,
            21,
            0,
            40,
            21,
            38,
            97,
            22,
            17,
            20,
            37,
            17,
            20,
        ],
        [
            42,
            20,
            18,
            0,
            40,
            57,
            43,
            42,
            0,
            0,
            19,
            17,
            0,
            39,
            0,
            0,
            22,
            41,
            0,
            0,
            17,
            42,
            40,
            43,
        ],
        [
            0,
            0,
            0,
            40,
            0,
            83,
            60,
            0,
            44,
            0,
            37,
            60,
            0,
            40,
            0,
            0,
            40,
            19,
            0,
            39,
            41,
            17,
            0,
            0,
        ],
        [
            60,
            18,
            0,
            37,
            18,
            0,
            21,
            41,
            23,
            39,
            0,
            63,
            60,
            39,
            0,
            19,
            0,
            16,
            16,
            0,
            0,
            40,
            0,
            16,
        ],
        [
            22,
            0,
            0,
            0,
            0,
            0,
            0,
            61,
            36,
            80,
            96,
            19,
            19,
            41,
            16,
            0,
            0,
            0,
            22,
            0,
            43,
            0,
            44,
            22,
        ],
        [
            0,
            0,
            20,
            19,
            17,
            20,
            40,
            0,
            0,
            60,
            61,
            0,
            20,
            62,
            20,
            0,
            0,
            38,
            0,
            0,
            0,
            0,
            24,
            22,
        ],
        [
            21,
            22,
            38,
            0,
            44,
            20,
            40,
            39,
            0,
            36,
            22,
            21,
            19,
            39,
            19,
            0,
            0,
            21,
            24,
            16,
            23,
            21,
            37,
            0,
        ],
        [
            0,
            24,
            23,
            39,
            20,
            0,
            0,
            41,
            0,
            0,
            0,
            22,
            0,
            0,
            44,
            42,
            22,
            42,
            22,
            19,
            20,
            58,
            18,
            0,
        ],
        [
            60,
            57,
            0,
            0,
            16,
            0,
            16,
            37,
            0,
            0,
            0,
            44,
            63,
            0,
            18,
            0,
            17,
            18,
            0,
            0,
            0,
            100,
            24,
            23,
        ],
        [
            0,
            44,
            44,
            0,
            23,
            17,
            39,
            21,
            0,
            17,
            40,
            0,
            24,
            78,
            17,
            24,
            20,
            18,
            0,
            24,
            24,
            0,
            20,
            0,
        ],
        [
            23,
            16,
            0,
            0,
            0,
            23,
            0,
            0,
            0,
            0,
            43,
            58,
            0,
            0,
            24,
            60,
            0,
            0,
            19,
            0,
            21,
            0,
            20,
            0,
        ],
        [
            44,
            20,
            0,
            19,
            21,
            0,
            39,
            19,
            0,
            0,
            0,
            39,
            22,
            0,
            0,
            64,
            24,
            22,
            0,
            39,
            0,
            43,
            42,
            16,
        ],
        [
            0,
            60,
            37,
            18,
            0,
            0,
            0,
            20,
            0,
            41,
            43,
            16,
            43,
            24,
            0,
            0,
            18,
            18,
            0,
            44,
            20,
            0,
            21,
            37,
        ],
        [
            0,
            0,
            23,
            39,
            0,
            24,
            40,
            0,
            37,
            0,
            40,
            20,
            44,
            43,
            0,
            0,
            0,
            0,
            0,
            16,
            0,
            59,
            0,
            0,
        ],
        [
            0,
            42,
            0,
            0,
            23,
            24,
            38,
            19,
            36,
            0,
            20,
            60,
            57,
            0,
            23,
            40,
            0,
            0,
            0,
            16,
            42,
            0,
            23,
            0,
        ],
        [
            0,
            41,
            36,
            43,
            23,
            41,
            17,
            0,
            38,
            0,
            0,
            21,
            21,
            17,
            16,
            16,
            39,
            0,
            22,
            0,
            21,
            23,
            16,
            23,
        ],
        [
            17,
            0,
            23,
            23,
            20,
            0,
            17,
            58,
            17,
            0,
            20,
            0,
            17,
            24,
            0,
            0,
            17,
            42,
            0,
            58,
            19,
            22,
            0,
            24,
        ],
        [
            42,
            0,
            16,
            0,
            43,
            0,
            24,
            36,
            0,
            16,
            24,
            41,
            41,
            0,
            24,
            0,
            0,
            0,
            0,
            0,
            56,
            38,
            63,
            19,
        ],
        [
            37,
            23,
            23,
            0,
            42,
            16,
            23,
            76,
            23,
            0,
            0,
            24,
            20,
            41,
            20,
            24,
            40,
            23,
            0,
            0,
            0,
            39,
            20,
            0,
        ],
        [
            43,
            20,
            17,
            17,
            0,
            20,
            19,
            0,
            80,
            0,
            0,
            0,
            40,
            0,
            40,
            16,
            19,
            0,
            0,
            0,
            18,
            0,
            17,
            0,
        ],
        [
            18,
            20,
            44,
            40,
            21,
            18,
            0,
            20,
            0,
            0,
            16,
            24,
            0,
            0,
            19,
            18,
            0,
            17,
            23,
            0,
            23,
            44,
            0,
            42,
        ],
        [
            44,
            0,
            0,
            0,
            62,
            0,
            17,
            41,
            0,
            0,
            0,
            63,
            0,
            37,
            22,
            0,
            20,
            0,
            0,
            19,
            57,
            18,
            38,
            0,
        ],
    ]

    for base in sorted(bases):
        # from base to obtain n, such like network05 → n=5
        try:
            n_str = "".join(ch for ch in base if ch.isdigit())
            n_val = int(n_str) if n_str else 5
        except Exception:
            print(f"[{base}] Cannot parse n from base, skip.")
            continue

        try:
            instance_data = {
                "n": n_val,
                "t": cut_matrix(your_t_0based, n_val),
                "M": 1000,
                "intscale": 1000,
            }

            interpreter = jm.Interpreter(instance_data)
            ommx_instance = interpreter.eval_problem(problem)

            sol_path = _pick_solution_file(sol_root, base)
            solution = None
            if sol_path:
                try:
                    print(f"  → Evaluating solution: {sol_path}")
                    solution_dict = parse_solution_zfx(sol_path, n_val)
                    solution = ommx_instance.evaluate(solution_dict)
                    if (
                        solution.feasible
                        and abs(solution.objective - solution_dict[0]) < 1e-6
                    ):
                        print(
                            f"    objective={solution.objective}, feasible={solution.feasible}"
                        )
                    else:
                        print(
                            "    ! Objective mismatch or infeasible; will save instance only."
                        )
                        solution = None
                except Exception as sol_err:
                    print(f"    ! Solution evaluation failed: {sol_err}")
                    solution = None

            out_path = os.path.join(output_directory, f"{base}.ommx")
            if os.path.exists(out_path):
                os.remove(out_path)

            # Add annotations to the instance.
            ommx_instance.title = base
            ommx_instance.license = LICENSE
            ommx_instance.dataset = "Network Design"
            ommx_instance.authors = QOBLIB_AUTHORS
            ommx_instance.num_variables = len(ommx_instance.decision_variables)
            ommx_instance.num_constraints = len(ommx_instance.constraints)
            ommx_instance.annotations["org.ommx.qoblib.url"] = (
                "https://git.zib.de/qopt/qoblib-quantum-optimization-benchmarking-library/-/tree/main/08-network?ref_type=heads"
            )

            builder = ArtifactBuilder.new_archive_unnamed(out_path)
            instance_desc = builder.add_instance(ommx_instance)
            if solution is not None:
                solution.instance = instance_desc.digest
                solution.annotations["org.ommx.qoblib.authors"] = QOBLIB_AUTHORS_STR
                builder.add_solution(solution)
            builder.build()

            print(f"  ✓ Created: {out_path}")
            print("-" * 50)
            processed_count += 1

        except Exception as e:
            print(f"[{base}] Error: {e}")
            print("-" * 50)
            error_count += 1

    print(f"Batch complete — processed: {processed_count}, errors: {error_count}")


if __name__ == "__main__":
    batch_process(
        sol_root="../../solutions",
        output_directory="./ommx_output",
    )
