# "commands/title/tmd.py" from WiiPy by NinjaCheetah
# https://github.com/NinjaCheetah/WiiPy

import pathlib
import libWiiPy
from modules.core import fatal_error
from modules.title import title_edit_ios, title_edit_tid, title_edit_type


def handle_tmd_edit(args):
    input_path = pathlib.Path(args.input)
    if args.output is not None:
        output_path = pathlib.Path(args.output)
    else:
        output_path = pathlib.Path(args.input)

    tmd = libWiiPy.title.TMD()
    tmd.load(input_path.read_bytes())

    # State variable to make sure that changes are made.
    edits_made = False
    # Go over every possible change, and apply them if they were specified.
    if args.tid is not None:
        new_tid = title_edit_tid(tmd.title_id, args.tid)
        tmd.set_title_id(new_tid)
        edits_made = True
    if args.ios is not None:
        new_ios_tid = title_edit_ios(args.ios)
        tmd.ios_tid = new_ios_tid
        edits_made = True
    if args.type is not None:
        new_tid = title_edit_type(tmd.title_id, args.type)
        tmd.set_title_id(new_tid)
        edits_made = True

    if not edits_made:
        fatal_error("You must specify at least one change to make!")

    # Fakesign the title since any changes have already invalidated the signature.
    tmd.fakesign()
    output_path.write_bytes(tmd.dump())

    print("Successfully edited TMD file!")


def handle_tmd_remove(args):
    input_path = pathlib.Path(args.input)
    if args.output is not None:
        output_path = pathlib.Path(args.output)
    else:
        output_path = pathlib.Path(args.input)

    if not input_path.exists():
        fatal_error("The specified TMD files does not exist!")

    tmd = libWiiPy.title.TMD()
    tmd.load(input_path.read_bytes())

    if args.index is not None:
        # Make sure the target index exists, then remove it from the TMD.
        if args.index >= len(tmd.content_records):
            fatal_error("The specified index could not be found in the provided TMD!")
        tmd.content_records.pop(args.index)
        tmd.num_contents -= 1
        # Auto fakesign because we've edited the TMD.
        tmd.fakesign()
        output_path.write_bytes(tmd.dump())
        print(f"Removed content record at index {args.index}!")

    elif args.cid is not None:
        if len(args.cid) != 8:
            fatal_error("The specified Content ID is invalid!")
        target_cid = int(args.cid, 16)
        # List Contents IDs in the title, and ensure that the target Content ID exists.
        valid_ids = []
        for record in tmd.content_records:
            valid_ids.append(record.content_id)
        if target_cid not in valid_ids:
            fatal_error("The specified Content ID could not be found in the provided TMD!")
        tmd.content_records.pop(valid_ids.index(target_cid))
        tmd.num_contents -= 1
        # Auto fakesign because we've edited the TMD.
        tmd.fakesign()
        output_path.write_bytes(tmd.dump())
        print(f"Removed content record with Content ID \"{target_cid:08X}\"!")