"""
Macro injection and extraction utilities.

Phase 3: Full VBA compilation and signing support
- VBA source compilation to binary format
- Authenticode signing for security
- Pure ZIP manipulation (no Office automation)
- Deterministic build process
"""

from __future__ import annotations

import tempfile
from pathlib import Path
from typing import Optional
import zipfile
import logging

from .vba_compiler import VbaCompiler
from .authenticode_signer import AuthenticodeSigner

logger = logging.getLogger(__name__)

SUPPORTED_EXTENSIONS = {
    ".dotx",
    ".docx",
    ".dotm",
    ".pptx",
    ".potx",
    ".potm",
    ".xlsx",
    ".xltx",
    ".xltm",
}

# Map file extension prefix to OOXML part directory
_EXT_TO_PART_DIR = {
    ".doc": "word",
    ".dot": "word",
    ".ppt": "ppt",
    ".pot": "ppt",
    ".xls": "xl",
    ".xlt": "xl",
}

_VBA_CONTENT_TYPE = "application/vnd.ms-office.vbaProject"
_VBA_REL_TYPE = "http://schemas.microsoft.com/office/2006/relationships/vbaProject"


def _get_part_dir(ext: str) -> str:
    """Get the OOXML part directory for a file extension."""
    prefix = ext[:4].lower()
    return _EXT_TO_PART_DIR.get(prefix, "word")


def extract_vba(artifact: Path, output_dir: Path) -> None:
    """Extract vbaProject.bin from an OOXML artifact.

    Writes raw binary for hashing/attestation.
    """
    if artifact.suffix.lower() not in SUPPORTED_EXTENSIONS:
        raise ValueError(f"Unsupported file type: {artifact}")

    part_dir = _get_part_dir(artifact.suffix.lower())
    vba_path = f"{part_dir}/vbaProject.bin"

    with zipfile.ZipFile(artifact, "r") as zf:
        names = zf.namelist()
        # Check both subdirectory and root for backwards compat
        actual_path = (
            vba_path
            if vba_path in names
            else ("vbaProject.bin" if "vbaProject.bin" in names else None)
        )
        if actual_path is None:
            return
        output_dir.mkdir(parents=True, exist_ok=True)
        data = zf.read(actual_path)
        (output_dir / "vbaProject.bin").write_bytes(data)


def inject_vba(template: Path, vba_binary: Path, output: Path) -> None:
    """Inject a vbaProject.bin into a macro-free template.

    This preserves all existing ZIP entries, adds the VBA binary at the
    correct OOXML part path, and updates [Content_Types].xml and the
    appropriate .rels file so Office recognizes the macro project.
    """
    if not vba_binary.exists():
        raise FileNotFoundError(vba_binary)

    from lxml import etree

    part_dir = _get_part_dir(template.suffix.lower())
    vba_part = f"{part_dir}/vbaProject.bin"
    rels_path = f"{part_dir}/_rels/{part_dir.split('/')[-1]}.xml.rels"
    # Word uses word/_rels/document.xml.rels, ppt uses ppt/_rels/presentation.xml.rels, xl uses xl/_rels/workbook.xml.rels
    rels_names = {
        "word": "word/_rels/document.xml.rels",
        "ppt": "ppt/_rels/presentation.xml.rels",
        "xl": "xl/_rels/workbook.xml.rels",
    }
    rels_path = rels_names.get(part_dir, f"{part_dir}/_rels/document.xml.rels")

    with zipfile.ZipFile(template, "r") as src:
        with zipfile.ZipFile(output, "w", zipfile.ZIP_DEFLATED) as dst:
            for info in src.infolist():
                data = src.read(info.filename)

                if info.filename == "[Content_Types].xml":
                    # Add VBA content type override
                    root = etree.fromstring(data)
                    ns = root.nsmap.get(
                        None,
                        "http://schemas.openxmlformats.org/package/2006/content-types",
                    )
                    # Check if override already exists
                    existing = [e for e in root if e.get("PartName") == f"/{vba_part}"]
                    if not existing:
                        override = etree.SubElement(root, f"{{{ns}}}Override")
                        override.set("PartName", f"/{vba_part}")
                        override.set("ContentType", _VBA_CONTENT_TYPE)
                    data = etree.tostring(
                        root, xml_declaration=True, encoding="UTF-8", standalone=True
                    )

                elif info.filename == rels_path:
                    # Add VBA relationship
                    root = etree.fromstring(data)
                    ns = root.nsmap.get(
                        None,
                        "http://schemas.openxmlformats.org/package/2006/relationships",
                    )
                    existing = [e for e in root if e.get("Type") == _VBA_REL_TYPE]
                    if not existing:
                        # Generate unique rId
                        used_ids = {e.get("Id", "") for e in root}
                        rid_num = 1
                        while f"rId{rid_num}" in used_ids:
                            rid_num += 1
                        rel = etree.SubElement(root, f"{{{ns}}}Relationship")
                        rel.set("Id", f"rId{rid_num}")
                        rel.set("Type", _VBA_REL_TYPE)
                        rel.set("Target", "vbaProject.bin")
                    data = etree.tostring(
                        root, xml_declaration=True, encoding="UTF-8", standalone=True
                    )

                dst.writestr(info, data)
            dst.writestr(vba_part, vba_binary.read_bytes())


def compile_and_inject_vba(
    template: Path,
    vba_source_dir: Path,
    output: Path,
    project_config: Optional[dict] = None,
    certificate_path: Optional[Path] = None,
    certificate_password: Optional[str] = None,
) -> bool:
    """Compile VBA source files and inject into template with optional signing.

    Args:
        template: Path to macro-free template (.dotx, .potx, etc.)
        vba_source_dir: Directory containing .bas, .cls, .frm files
        output: Where to save the final document with macros
        project_config: VBA project configuration
        certificate_path: Path to code signing certificate (.pfx/.p12)
        certificate_password: Certificate password

    Returns:
        True if compilation and injection successful
    """
    try:
        # Step 1: Compile VBA source to binary
        compiler = VbaCompiler()
        with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as tmp_file:
            vba_binary_path = Path(tmp_file.name)

        success = compiler.compile_project(
            vba_source_dir, vba_binary_path, project_config
        )
        if not success:
            logger.error("VBA compilation failed")
            vba_binary_path.unlink(missing_ok=True)
            return False

        # Step 2: Sign the compiled VBA project (optional)
        if certificate_path:
            signer = AuthenticodeSigner()
            signing_success = signer.sign_vba_project(
                vba_binary_path, certificate_path, certificate_password
            )
            if not signing_success:
                logger.warning(
                    "VBA project signing failed, proceeding without signature"
                )

        # Step 3: Inject into template
        inject_vba(template, vba_binary_path, output)

        # Cleanup
        vba_binary_path.unlink(missing_ok=True)

        logger.info(f"Successfully compiled and injected VBA macros into: {output}")
        return True

    except Exception as e:
        logger.error(f"VBA compilation and injection failed: {e}")
        return False


def sign_office_document(
    document_path: Path,
    certificate_path: Path,
    certificate_password: Optional[str] = None,
    timestamp_url: Optional[str] = None,
) -> bool:
    """Sign an Office document with Authenticode certificate.

    Args:
        document_path: Path to Office document (.docx, .xlsx, .pptx)
        certificate_path: Path to .pfx/.p12 certificate file
        certificate_password: Certificate password
        timestamp_url: Timestamp server URL

    Returns:
        True if signing successful
    """
    signer = AuthenticodeSigner()
    return signer.sign_office_document(
        document_path, certificate_path, certificate_password, timestamp_url
    )
