#@title Headless (Qt off) + dependencies
import os
os.environ["QT_QPA_PLATFORM"] = "offscreen"
os.environ["MPLBACKEND"] = "Agg"

import matplotlib as mpl
try:
    mpl.use("Agg", force=True)
except TypeError:
    mpl.use("Agg")
print("Matplotlib backend:", mpl.get_backend())

# Scientific + GDS toolchain (incl. Descartes)
!pip install "jedi>=0.16"
%pip -q install --upgrade pip wheel setuptools
%pip -q install "numpy>=1.24" "matplotlib>=3.8" \
                "gdstk>=0.9.61" "shapely>=2.0" "ezdxf>=1.2.0" \
                "pandas>=2.0" "scipy>=1.10" "networkx>=2.8" \
                "pint>=0.20" "addict>=2.4.0" "pyyaml>=6.0.1" \
                "qutip>=4.7" "h5py>=3.8" "descartes>=1.1" "jedi>=0.19.1"

# @title Clone Qiskit Metal; Headless Setup
"""Setup script for Qiskit Metal in a headless Google Colab environment.

This script clones the Qiskit Metal repository, patches the initialization 
files to remove GUI dependencies (PySide2), and configures the Python path.
"""

import logging
import os
import re
import sys
import textwrap
from pathlib import Path

# Constants
REPO_URL = "https://github.com/qiskit-community/qiskit-metal"
ROOT_DIR = Path("/content/qiskit-metal")
SRC_DIR = ROOT_DIR / "src"
PKG_DIR = SRC_DIR / "qiskit_metal"

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def setup_repository():
    """Clones the Qiskit Metal repository and verifies the structure.

    Returns:
        Path: The path to the package directory.

    Raises:
        RuntimeError: If the package directory is not found after cloning.
    """
    logger.info("Cleaning old installation and cloning repository...")
    !rm -rf {ROOT_DIR}
    !git clone --depth 1 {REPO_URL} {ROOT_DIR}

    if not PKG_DIR.exists():
        logger.error("Package folder not found at %s", PKG_DIR)
        if ROOT_DIR.exists():
            logger.info("Root directory contents: %s", os.listdir(ROOT_DIR))
        raise RuntimeError(
            f"Git clone failed to provide the qiskit_metal package folder "
            f"at {PKG_DIR}. The repo structure may have changed."
        )

    return PKG_DIR


def configure_python_path():
    """Configures sys.path and PYTHONPATH to include the src directory."""
    if str(SRC_DIR) not in sys.path:
        sys.path.insert(0, str(SRC_DIR))
    
    current_pythonpath = os.environ.get("PYTHONPATH", "")
    os.environ["PYTHONPATH"] = f"{SRC_DIR}:{current_pythonpath}"
    logger.info("Python path configured to include %s", SRC_DIR)


def patch_package_init():
    """Replaces the package __init__.py with a minimal headless version."""
    init_file = PKG_DIR / "__init__.py"
    if not init_file.exists():
        return

    # Backup original
    orig_init = init_file.read_text(encoding="utf-8")
    (PKG_DIR / "__init__orig.py").write_text(orig_init, encoding="utf-8")

    minimal_init = textwrap.dedent("""
        # [colab] Headless, layout-only __init__
        import logging as _logging
        try:
            from addict import Dict as Dict
        except Exception:
            try:
                from .toolbox_python.attr_dict import Dict
            except Exception:
                class Dict(dict): pass

        logger = _logging.getLogger("qiskit_metal_colab")

        class _Config:
            @staticmethod
            def is_building_docs():
                return False
        config = _Config()

        def is_design(obj):
            try:
                from .designs.design_base import QDesign
                return isinstance(obj, QDesign)
            except Exception:
                return False

        def is_component(obj):
            try:
                from .qlibrary.core.base import QComponent
                return isinstance(obj, QComponent)
            except Exception:
                return False

        __all__ = ["Dict", "config", "logger", "is_design", "is_component"]
    """).strip() + "\n"
    
    init_file.write_text(minimal_init, encoding="utf-8")
    logger.info("Patched %s for headless mode.", init_file)


def scrub_draw_imports():
    """Removes PySide2/Qt dependencies from the draw module initialization."""
    draw_init = PKG_DIR / "draw" / "__init__.py"
    if not draw_init.exists():
        return

    content = draw_init.read_text(encoding="utf-8")

    # Guard "from . import mpl"
    content = re.sub(
        r'^\s*from\s+\.\s*import\s+mpl\s*$',
        "try:\n    from . import mpl\n"
        "except Exception as _e:\n"
        "    print('[colab] draw.mpl disabled (headless):', _e)\n",
        content,
        flags=re.MULTILINE,
    )

    # Guard "from .mpl import ..."
    content = re.sub(
        r'^\s*from\s+\.\s*mpl\s+import[^\n]*$',
        "try:\n    from .mpl import render, figure_spawn\n"
        "except Exception as _e:\n"
        "    print('[colab] draw.mpl (named) disabled (headless):', _e)\n"
        "    def render(*a, **k):\n"
        "        raise RuntimeError('draw.mpl unavailable in headless mode')\n"
        "    def figure_spawn(*a, **k):\n"
        "        raise RuntimeError('draw.mpl unavailable in headless mode')\n",
        content,
        flags=re.MULTILINE,
    )
    
    draw_init.write_text(content, encoding="utf-8")
    logger.info("Scrubbed Qt imports from %s", draw_init)


def disable_renderers():
    """Disables the renderers package to prevent Qt imports."""
    rndr_init = PKG_DIR / "renderers" / "__init__.py"
    if rndr_init.exists():
        # Backup
        (rndr_init.parent / "__init__orig.py").write_text(
            rndr_init.read_text(encoding="utf-8"), encoding="utf-8"
        )
        # Override
        rndr_init.write_text(
            "# [colab] minimal renderers package\n__all__ = []\n",
            encoding="utf-8",
        )
        logger.info("Disabled renderers package at %s", rndr_init)


def main():
    """Main execution flow for Qiskit Metal setup."""
    try:
        setup_repository()
        configure_python_path()
        patch_package_init()
        scrub_draw_imports()
        disable_renderers()

        # Final Verification
        import importlib
        importlib.invalidate_caches()
        
        import qiskit_metal
        print(f"Success! qiskit_metal imported from: {qiskit_metal.__file__}")

    except Exception as e:
        logger.error("Setup failed: %s", e)
        raise


if __name__ == "__main__":
    main()