"""
Switch Downgrader — Downgrade Engine.
Contains the Quick Downgrade (Level 1) and Full Downgrade (Level 2) processes.
"""
import os
import re
import shutil
import stat
import tempfile
import datetime
from pathlib import Path

from core.utils import (
    run_command, run_command_with_progress, run_and_interrupt_flash,
    safe_remove_directory, check_disk_space, get_base_path
)
from core.detection import detect_console_model, detect_switch_drives_wmi
from core.i18n import t


class DowngradeEngine:
    """Core downgrade engine — stateless process runner."""

    def __init__(self, paths, log_func, progress_func, dialog_func, offline_mode=False):
        """
        paths: dict of StringVar-like objects with .get() → str
        log_func: callable(message)
        progress_func: callable(progress_text)
        dialog_func: callable(title, message, buttons="ok") → bool
        offline_mode: bool
        """
        self.paths = paths
        self.log = log_func
        self.progress = progress_func
        self.dialog = dialog_func
        self.offline = offline_mode
        self.last_output_dir = None

    # ------------------------------------------------------------------
    # Shared helpers
    # ------------------------------------------------------------------

    def _cmd(self, command, cwd=None):
        return run_command(command, cwd=cwd, log_func=self.log)

    def _cmd_progress(self, command, task_name="Processing"):
        return run_command_with_progress(command, task_name, log_func=self.log, progress_func=self.progress)

    def _get_nand_source(self):
        """Resolve NAND source (offline file or online eMMC device).
        Returns (path, type_str) or (None, None)."""
        if self.offline:
            rawnand = self.paths['rawnand'].get()
            if not rawnand or not Path(rawnand).exists():
                self.log("ERROR: RAWNAND.bin file not found.")
                return None, None
            self.log(t("process.connecting_nand", path=rawnand))
            return rawnand, 'file'

        # Online — detect physical eMMC
        self.log(t("process.detecting_emmc"))
        self.log(t("process.detecting_emmc_hint"))

        drives = detect_switch_drives_wmi(log_func=self.log)
        if not drives:
            self.dialog(t("emmc_dialog.not_found_title"), t("emmc_dialog.not_found_msg"))
            return None, None

        if len(drives) > 1:
            self.dialog(t("emmc_dialog.multiple_title"), t("emmc_dialog.multiple_msg"))
            return None, None

        target = drives[0]
        msg = t("emmc_dialog.confirm_msg",
                type=target.get('type', 'eMMC'),
                path=target['path'],
                size=target['size'],
                model=target['model'])
        if not self.dialog(t("emmc_dialog.confirm_title"), msg, "yesno"):
            self.log("--- User cancelled.")
            return None, None

        self.log(f"SUCCESS: User confirmed eMMC at {target['path']}")
        self._last_drive_type = target.get('type', 'eMMC GPP')
        return target['path'], 'device'

    def _dump_prodinfo(self, nand_source, source_type, temp_dir):
        """Dump and validate PRODINFO. Returns (prodinfo_path, model_name, is_mariko) or (None,None,None)."""
        nx = self.paths['nxnandmanager'].get()
        keys = self.paths['keys'].get()
        prodinfo_path = Path(temp_dir) / "PRODINFO"

        self.log(f"Dumping PRODINFO from {source_type}...")
        dump_cmd = [nx, '-i', nand_source, '-keyset', keys, '-o', temp_dir, '-d', '-part=PRODINFO']
        rc, _ = self._cmd(dump_cmd)

        if rc != 0 or not prodinfo_path.exists():
            self.log(f"ERROR: Could not dump PRODINFO from {source_type}.")
            self.dialog(t("errors.prodinfo_title"), t("errors.prodinfo_dump_fail", source=source_type))
            return None, None, None

        # Validate magic
        with open(prodinfo_path, 'rb') as f:
            if f.read(4) != b'CAL0':
                self.log("ERROR: PRODINFO is invalid or encrypted.")
                self.dialog(t("errors.prodinfo_title"), t("errors.prodinfo_invalid"))
                return None, None, None

        self.log("SUCCESS: PRODINFO is valid and decrypted.")
        model_name, is_mariko = detect_console_model(str(prodinfo_path), log_func=self.log)
        if model_name is None:
            self.log("WARNING: Could not detect console model. Assuming Mariko.")
            model_name, is_mariko = "Unknown (Mariko assumed)", True

        return str(prodinfo_path), model_name, is_mariko

    def _generate_boot_files(self, temp_dir, is_mariko):
        """Run EmmcHaccGen to generate boot/system files. Returns output_dir Path or None."""
        keys = self.paths['keys'].get()
        fw = self.paths['firmware'].get()
        emmchaccgen = self.paths['emmchaccgen'].get()

        out_dir = Path(temp_dir) / "emmchaccgen_out"
        out_dir.mkdir(exist_ok=True)

        cmd = [emmchaccgen, '--keys', keys, '--fw', fw]
        if is_mariko:
            self.log("--- Mariko-type console detected. Using --mariko flag.")
            cmd.append('--mariko')
        else:
            self.log("--- Erista console detected. Using --no-autorcm flag.")
            cmd.append('--no-autorcm')

        rc, _ = self._cmd(cmd, cwd=str(out_dir))
        if rc != 0:
            self.log("ERROR: EmmcHaccGen failed to generate boot files.")
            self.dialog(t("errors.generation_title"), t("errors.generation_fail"))
            return None
        return out_dir

    def _get_versioned_folder(self, emmchaccgen_out_dir):
        """Get the first versioned subfolder from EmmcHaccGen output."""
        return next(d for d in Path(emmchaccgen_out_dir).iterdir() if d.is_dir())

    def _flash_bcpkg2(self, versioned_folder, nand_source):
        """Flash all 4 BCPKG2 partitions."""
        nx = self.paths['nxnandmanager'].get()
        parts = ["BCPKG2-1-Normal-Main", "BCPKG2-2-Normal-Sub",
                 "BCPKG2-3-SafeMode-Main", "BCPKG2-4-SafeMode-Sub"]
        for name in parts:
            f = versioned_folder / f"{name}.bin"
            if not f.exists():
                self.log(f"ERROR: {f.name} not found.")
                return False
            cmd = [nx, '-i', str(f), '-o', nand_source, f'-part={name}', 'FORCE']
            if self._cmd(cmd)[0] != 0:
                self.log(f"ERROR: Failed to flash {name}.")
                return False
        self.log("SUCCESS: All BCPKG2 partitions flashed.")
        return True

    def _save_boot_files(self, versioned_folder, dest_dir):
        """Copy BOOT0.bin and BOOT1.bin to destination."""
        shutil.copy2(versioned_folder / "BOOT0.bin", Path(dest_dir) / "BOOT0")
        shutil.copy2(versioned_folder / "BOOT1.bin", Path(dest_dir) / "BOOT1")
        self.log(f"SUCCESS: BOOT0 and BOOT1 saved to {dest_dir}")

    def _selective_copy_system_level1(self, source_system_path, drive_letter):
        """Level 1: Delete registered folder, then merge EmmcHaccGen SYSTEM output."""
        try:
            self.log("--- Updating system partition...")
            registered = drive_letter / "Contents" / "registered"
            if registered.exists():
                self.log("--- Removing existing registered folder...")
                if not safe_remove_directory(registered):
                    self.log("ERROR: Could not remove registered folder")
                    return False

            def merge_copy(src, dst):
                for item in src.iterdir():
                    dst_item = dst / item.name
                    if item.is_dir():
                        dst_item.mkdir(exist_ok=True)
                        merge_copy(item, dst_item)
                    else:
                        shutil.copy2(item, dst_item)

            self.log("--- Copying files from EmmcHaccGen output...")
            merge_copy(source_system_path, drive_letter)
            self.log("--- SUCCESS: System partition updated.")
            return True
        except Exception as e:
            self.log(f"ERROR: Failed to update system partition: {e}")
            return False

    def run_quick_downgrade(self, temp_dir):
        """Quick Downgrade: fixes SYSTEM + BCPKG2 for firmware downgrade.
        Preserves user data."""
        mode_label = "OFFLINE" if self.offline else "ONLINE"
        self.log(t("quick.title", mode=mode_label))

        if not self.offline:
            self.log(t("quick.desc_online"))
        else:
            self.log(t("quick.desc_offline"))

        if not check_disk_space(60, log_func=self.log, temp_path=temp_dir):
            return False

        # Step 1 — NAND source
        self.log(t("process.step_nand", n=1, total=6))
        nand_source, source_type = self._get_nand_source()
        if not nand_source:
            return False

        nx = self.paths['nxnandmanager'].get()
        keys = self.paths['keys'].get()

        self.log(t("process.step_prodinfo", n=2, total=6))
        prodinfo_path, model_name, is_mariko = self._dump_prodinfo(nand_source, source_type, temp_dir)
        if not prodinfo_path:
            return False
        self.log(t("process.step_model", model=model_name))

        self.log(t("process.step_generate", n=3, total=6))
        out_dir = self._generate_boot_files(temp_dir, is_mariko)
        if not out_dir:
            return False

        self.log(t("process.step_dump_system", n=4, total=6))
        dump_cmd = [nx, '-i', nand_source, '-keyset', keys, '-o', temp_dir, '-d', '-part=SYSTEM']
        if self._cmd(dump_cmd)[0] != 0:
            return False

        system_path = Path(temp_dir) / "SYSTEM"
        if not system_path.exists():
            self.log("ERROR: SYSTEM file was not created.")
            return False

        self.log("--- Mounting SYSTEM partition...")
        osfmount = self.paths['osfmount'].get()
        rc, output = self._cmd([osfmount, '-a', '-t', 'file', '-f', str(system_path), '-o', 'rw', '-m', '#:'])
        if rc != 0:
            self.dialog(t("errors.mount_title"), t("errors.mount_fail"))
            return False

        match = re.search(r"([A-Z]:)", output)
        if not match:
            self.log("ERROR: Could not determine drive letter.")
            return False

        drive_str = match.group(1)
        drive = Path(drive_str + "\\")
        self.log(f"   Mounted at {drive_str}")

        try:
            versioned = self._get_versioned_folder(out_dir)
            source_sys = versioned / "SYSTEM"
            if not self._selective_copy_system_level1(source_sys, drive):
                return False
        finally:
            self.log("--- Dismounting SYSTEM...")
            self._cmd([osfmount, '-D', '-m', drive_str])

        self.log(t("process.step_modified", n=5, total=6))
        flash_cmd = [nx, '-i', str(system_path), '-o', nand_source, '-part=SYSTEM', '-e', '-keyset', keys, 'FORCE']
        if self._cmd(flash_cmd)[0] != 0:
            self.log("ERROR: Failed to flash SYSTEM partition.")
            return False
        self.log("SUCCESS: SYSTEM partition restored.")

        self.log(t("process.step_flash_boot", n=6, total=6))
        versioned = self._get_versioned_folder(out_dir)
        if not self._flash_bcpkg2(versioned, nand_source):
            return False

        if self.offline:
            rawnand_folder = str(Path(nand_source).parent)
            self._save_boot_files(versioned, rawnand_folder)
            self.log(t("quick.complete"))
            self.dialog(t("ui.title"), t("quick.offline_done_msg"))
        else:
            self._save_boot_files(versioned, temp_dir)
            self.last_output_dir = temp_dir
            self.log(t("quick.complete"))
            self.log(t("quick.online_next_steps"))
            self.dialog(t("ui.title"), t("quick.online_done_msg"))

        return True

    # ------------------------------------------------------------------
