#!/usr/bin/env python3
"""
Ubuntu ISO Builder for Autoinstall
Follows the Ubuntu autoinstall guide exactly:
1. Download Ubuntu ISO
2. Extract with 7zip equivalent
3. Modify GRUB configuration
4. Add server/ directory with autoinstall config
5. Rebuild ISO with xorriso equivalent
"""

import os
import logging
import subprocess
import tempfile
import shutil
import urllib.request
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# Ubuntu ISO URLs with fallbacks
UBUNTU_MIRRORS = [
    "https://releases.ubuntu.com",  # Official Ubuntu releases (primary)
    "http://archive.ubuntu.com/ubuntu-releases",  # Official archive
    "http://mirror.csclub.uwaterloo.ca/ubuntu-releases",  # University of Waterloo
    "http://mirror.math.princeton.edu/pub/ubuntu-releases",  # Princeton mirror
]
UBUNTU_24_04_ISO = "24.04/ubuntu-24.04.3-live-server-amd64.iso"
UBUNTU_22_04_ISO = "22.04/ubuntu-22.04.3-live-server-amd64.iso"


def get_latest_ubuntu_server_iso(version: str) -> Optional[str]:
    """Dynamically find the latest point release for a given Ubuntu version."""
    try:
        import re

        for mirror in UBUNTU_MIRRORS:
            try:
                url = f"{mirror}/{version}/"
                response = urllib.request.urlopen(url, timeout=10)
                content = response.read().decode("utf-8")

                # Find all server ISO files for this version
                pattern = rf"ubuntu-{
                    re.escape(version)}\.[0-9]+-live-server-amd64\.iso"
                matches = re.findall(pattern, content)

                if matches:
                    # Sort by version and return the latest
                    matches.sort(key=lambda x: [int(n) for n in re.findall(r"\d+", x)])
                    latest_iso = matches[-1]
                    return f"{version}/{latest_iso}"
            except Exception as e:
                logger.debug(f"Failed to check {mirror}/{version}/: {e}")
                continue

        logger.warning(f"Could not dynamically detect latest ISO for Ubuntu {version}")
        return None

    except Exception as e:
        logger.error(f"Error in dynamic ISO detection: {e}")
        return None


class UbuntuISOBuilder:
    """Builds modified Ubuntu ISO with autoinstall configuration."""

    def __init__(self, vm_name: str, work_dir: Optional[str] = None):
        self.vm_name = vm_name
        self.work_dir = (
            Path(work_dir)
            if work_dir
            else Path(tempfile.mkdtemp(prefix="ubuntu-autoinstall-"))
        )
        self.source_files_dir = self.work_dir / "source-files"
        self.boot_dir = self.work_dir / "BOOT"
        self.server_dir = self.source_files_dir / "server"
        self.grub_cfg_path = self.source_files_dir / "boot" / "grub" / "grub.cfg"

        # Ensure directories exist
        self.work_dir.mkdir(exist_ok=True, parents=True)
        self.source_files_dir.mkdir(exist_ok=True, parents=True)

    def check_tools(self) -> bool:
        """Check if required tools are available."""

        # Check for 7zip equivalent (p7zip on macOS/Linux)
        if not shutil.which("7z") and not shutil.which("7za"):
            logger.error(
                "7zip not found. Install with: brew install p7zip (macOS) or apt install p7zip-full (Ubuntu)"
            )
            return False

        # Check for xorriso equivalent
        if (
            not shutil.which("xorriso")
            and not shutil.which("mkisofs")
            and not shutil.which("hdiutil")
        ):
            logger.error(
                "No ISO creation tool found. Install xorriso, mkisofs, or use macOS hdiutil"
            )
            return False

        return True

    def download_ubuntu_iso(self, version: str = "24.04") -> Path:
        """Download Ubuntu ISO if not already present, trying multiple mirrors."""
        iso_filename = f"ubuntu-{version}-live-server-amd64.iso"
        iso_path = self.work_dir / iso_filename

        if iso_path.exists():
            logger.info(f"Ubuntu ISO already exists: {iso_path}")
            return iso_path

        if version == "24.04":
            iso_subpath = UBUNTU_24_04_ISO
        elif version == "22.04":
            iso_subpath = UBUNTU_22_04_ISO
        else:
            raise ValueError(f"Unsupported Ubuntu version: {version}")

        # Try each mirror until one works
        last_error = None
        for mirror in UBUNTU_MIRRORS:
            iso_url = f"{mirror}/{iso_subpath}"
            logger.info(f"Trying to download Ubuntu {version} ISO from {iso_url}")

            try:
                # Try downloading from this mirror
                urllib.request.urlretrieve(iso_url, iso_path)
                logger.info(
                    f"✅ Ubuntu ISO downloaded successfully from {mirror}: {iso_path}"
                )
                return iso_path
            except Exception as e:
                last_error = e
                logger.warning(f"Failed to download from {mirror}: {e}")
                # Remove partial download if it exists
                if iso_path.exists():
                    iso_path.unlink()
                continue

        # If we get here, all mirrors failed
        logger.error(
            f"Failed to download Ubuntu ISO from all mirrors. Last error: {last_error}"
        )
        raise last_error

    def extract_iso(self, iso_path: Path) -> bool:
        """Extract Ubuntu ISO following the guide."""
        logger.info(f"Extracting ISO: {iso_path}")

        # Use 7z to extract ISO
        seven_zip_cmd = "7z" if shutil.which("7z") else "7za"

        try:
            # Extract ISO: 7z -y x ubuntu.iso -osource-files
            subprocess.run(
                [
                    seven_zip_cmd,
                    "-y",
                    "x",
                    str(iso_path),
                    f"-o{self.source_files_dir}",
                ],
                capture_output=True,
                text=True,
                check=True,
            )

            logger.info("ISO extracted successfully")

            # Move [BOOT] directory as per guide: mv '[BOOT]' ../BOOT
            boot_source = self.source_files_dir / "[BOOT]"
            if boot_source.exists():
                shutil.move(str(boot_source), str(self.boot_dir))
                logger.info(f"Moved [BOOT] directory to {self.boot_dir}")
            else:
                logger.warning("[BOOT] directory not found in extracted files")

            return True

        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to extract ISO: {e.stderr}")
            return False
        except Exception as e:
            logger.error(f"Error extracting ISO: {e}")
            return False

    def modify_grub_config(self) -> bool:
        """Modify GRUB configuration to add autoinstall menu entry."""
        logger.info("Modifying GRUB configuration...")

        if not self.grub_cfg_path.exists():
            logger.error(f"GRUB config not found: {self.grub_cfg_path}")
            return False

        try:
            # Read existing GRUB config
            with open(self.grub_cfg_path, "r", encoding="utf-8") as f:
                grub_content = f.read()

            # Autoinstall menu entry as per guide
            autoinstall_entry = """menuentry "Autoinstall Ubuntu Server" {
    set gfxpayload=keep
    linux   /casper/vmlinuz quiet autoinstall ds=nocloud\\;s=/cdrom/server/  ---
    initrd  /casper/initrd
}

"""

            # Insert autoinstall entry at the beginning of menu entries
            # Find the first menuentry and insert before it
            import re

            first_menu_match = re.search(r'(menuentry\s+["\'])', grub_content)
            if first_menu_match:
                insert_pos = first_menu_match.start()
                modified_content = (
                    grub_content[:insert_pos]
                    + autoinstall_entry
                    + grub_content[insert_pos:]
                )
            else:
                # Fallback: append at the end
                modified_content = grub_content + "\n" + autoinstall_entry

            # Write modified GRUB config
            with open(self.grub_cfg_path, "w", encoding="utf-8") as f:
                f.write(modified_content)

            logger.info("GRUB configuration modified successfully")
            return True

        except Exception as e:
            logger.error(f"Failed to modify GRUB config: {e}")
            return False

    def create_autoinstall_config(self, user_data: str) -> bool:
        """Create autoinstall configuration in server/ directory."""
        logger.info("Creating autoinstall configuration...")

        try:
            # Create server directory
            self.server_dir.mkdir(exist_ok=True, parents=True)

            # Create empty meta-data file (as per guide)
            meta_data_path = self.server_dir / "meta-data"
            meta_data_path.touch()
            logger.info(f"Created empty meta-data: {meta_data_path}")

            # Create user-data file with autoinstall configuration
            user_data_path = self.server_dir / "user-data"
            with open(user_data_path, "w", encoding="utf-8") as f:
                f.write(user_data)
            logger.info(f"Created user-data: {user_data_path}")

            return True

        except Exception as e:
            logger.error(f"Failed to create autoinstall config: {e}")
            return False

    def rebuild_iso(self, output_path: Path) -> bool:
        """Rebuild ISO with autoinstall configuration using xorriso."""
        logger.info(f"Rebuilding ISO: {output_path}")

        try:
            # Change to source-files directory for xorriso command
            original_cwd = os.getcwd()
            os.chdir(self.source_files_dir)

            # Remove existing output file
            if output_path.exists():
                output_path.unlink()

            # Try different ISO creation methods in order of preference
            success = False

            # Method 1: xorriso (most complete)
            if shutil.which("xorriso") and not success:
                try:
                    logger.info("Trying xorriso method...")
                    cmd = [
                        "xorriso",
                        "-as",
                        "mkisofs",
                        "-r",
                        "-V",
                        f"Ubuntu 24.04 LTS AUTO (EFIBIOS)",
                        "-o",
                        str(output_path),
                        "--grub2-mbr",
                        f"..{os.sep}BOOT{os.sep}1-Boot-NoEmul.img",
                        "-partition_offset",
                        "16",
                        "--mbr-force-bootable",
                        "-append_partition",
                        "2",
                        "28732ac11ff8d211ba4b00a0c93ec93b",
                        f"..{os.sep}BOOT{os.sep}2-Boot-NoEmul.img",
                        "-appended_part_as_gpt",
                        "-iso_mbr_part_type",
                        "a2a0d0ebe5b9334487c068b6b72699c7",
                        "-c",
                        "/boot.catalog",
                        "-b",
                        "/boot/grub/i386-pc/eltorito.img",
                        "-no-emul-boot",
                        "-boot-load-size",
                        "4",
                        "-boot-info-table",
                        "--grub2-boot-info",
                        "-eltorito-alt-boot",
                        "-e",
                        "--interval:appended_partition_2:::",
                        "-no-emul-boot",
                        ".",
                    ]
                    subprocess.run(cmd, capture_output=True, text=True, check=True)
                    success = True
                    logger.info("✅ ISO created with xorriso")
                except subprocess.CalledProcessError as e:
                    logger.warning(f"xorriso failed: {e.stderr}")
                    if output_path.exists():
                        output_path.unlink()

            # Method 2: mkisofs with joliet-long
            if shutil.which("mkisofs") and not success:
                try:
                    logger.info("Trying mkisofs with joliet-long...")
                    cmd = [
                        "mkisofs",
                        "-r",
                        "-V",
                        f"Ubuntu 24.04 LTS AUTO",
                        "-cache-inodes",
                        "-J",
                        "-joliet-long",
                        "-l",
                        "-b",
                        "boot/grub/i386-pc/eltorito.img",
                        "-c",
                        "boot.catalog",
                        "-no-emul-boot",
                        "-boot-load-size",
                        "4",
                        "-boot-info-table",
                        "-o",
                        str(output_path),
                        ".",
                    ]
                    subprocess.run(cmd, capture_output=True, text=True, check=True)
                    success = True
                    logger.info("✅ ISO created with mkisofs (joliet-long)")
                except subprocess.CalledProcessError as e:
                    logger.warning(f"mkisofs with joliet-long failed: {e.stderr}")
                    if output_path.exists():
                        output_path.unlink()

            # Method 3: mkisofs without Joliet (fallback)
            if shutil.which("mkisofs") and not success:
                try:
                    logger.info("Trying mkisofs without Joliet (fallback)...")
                    cmd = [
                        "mkisofs",
                        "-r",
                        "-V",
                        f"Ubuntu 24.04 LTS AUTO",
                        "-cache-inodes",
                        "-l",  # No -J (Joliet) to avoid filename conflicts
                        "-b",
                        "boot/grub/i386-pc/eltorito.img",
                        "-c",
                        "boot.catalog",
                        "-no-emul-boot",
                        "-boot-load-size",
                        "4",
                        "-boot-info-table",
                        "-o",
                        str(output_path),
                        ".",
                    ]
                    subprocess.run(cmd, capture_output=True, text=True, check=True)
                    success = True
                    logger.info("✅ ISO created with mkisofs (no Joliet)")
                except subprocess.CalledProcessError as e:
                    logger.warning(
                        f"mkisofs without Joliet failed: {
                            e.stderr}"
                    )
                    if output_path.exists():
                        output_path.unlink()

            # Method 4: macOS hdiutil
            if shutil.which("hdiutil") and not success:
                try:
                    logger.info("Trying hdiutil (macOS)...")
                    cmd = [
                        "hdiutil",
                        "makehybrid",
                        "-iso",
                        "-joliet",
                        "-o",
                        str(output_path),
                        ".",
                    ]
                    subprocess.run(cmd, capture_output=True, text=True, check=True)
                    success = True
                    logger.info("✅ ISO created with hdiutil")
                except subprocess.CalledProcessError as e:
                    logger.warning(f"hdiutil failed: {e.stderr}")
                    if output_path.exists():
                        output_path.unlink()

            if not success:
                logger.error("All ISO creation methods failed")
                return False

            # Verify the output file was created
            if not output_path.exists():
                logger.error("ISO file was not created despite success message")
                return False

            logger.info(f"ISO rebuilt successfully: {output_path}")
            logger.info(
                f"ISO size: {output_path.stat().st_size / (1024 * 1024):.1f} MB"
            )
            return True

        except Exception as e:
            logger.error(f"Error rebuilding ISO: {e}")
            return False
        finally:
            # Return to original directory
            os.chdir(original_cwd)

    def build_autoinstall_iso(
        self, user_data: str, output_path: Path, ubuntu_version: str = "24.04"
    ) -> bool:
        """Complete ISO build process following the Ubuntu autoinstall guide."""
        logger.info(
            f"🚀 Starting Ubuntu {ubuntu_version} autoinstall ISO build process"
        )

        try:
            # Step 1: Check tools
            if not self.check_tools():
                return False

            # Step 2: Download Ubuntu ISO
            iso_path = self.download_ubuntu_iso(ubuntu_version)

            # Step 3: Extract ISO
            if not self.extract_iso(iso_path):
                return False

            # Step 4: Modify GRUB
            if not self.modify_grub_config():
                return False

            # Step 5: Create autoinstall config
            if not self.create_autoinstall_config(user_data):
                return False

            # Step 6: Rebuild ISO
            if not self.rebuild_iso(output_path):
                return False

            logger.info(f"🎉 Successfully created autoinstall ISO: {output_path}")
            logger.info(f"📁 Work directory: {self.work_dir}")
            return True

        except Exception as e:
            logger.error(f"Failed to build autoinstall ISO: {e}")
            return False

    def cleanup(self):
        """Clean up temporary work directory."""
        if self.work_dir.exists():
            shutil.rmtree(self.work_dir)
            logger.info(f"Cleaned up work directory: {self.work_dir}")


def main():
    """Test the ISO builder."""
    import logging

    logging.basicConfig(level=logging.INFO)

    # Sample autoinstall user-data
    user_data = """#cloud-config
autoinstall:
  version: 1
  packages:
    - ubuntu-server
  identity:
    realname: 'Test User'
    username: testuser
    password: '$6$rounds=4096$saltsalt$[AWS-SECRET-REMOVED]AzpI8g8T14F8VnhXo0sUkZV2NV6/.c77tHgVi34DgbPu.'
    hostname: test-vm
  locale: en_US.UTF-8
  keyboard:
    layout: us
  storage:
    layout:
      name: direct
  ssh:
    install-server: true
  late-commands:
    - curtin in-target -- apt-get autoremove -y
"""

    builder = UbuntuISOBuilder("test-vm")
    output_path = Path("/tmp/ubuntu-24.04-autoinstall.iso")

    success = builder.build_autoinstall_iso(user_data, output_path)
    if success:
        print(f"✅ ISO created: {output_path}")
    else:
        print("❌ ISO creation failed")

    # Optionally clean up
    # builder.cleanup()


if __name__ == "__main__":
    main()