import json
import time
import os

import network

try:
    import mip
except ImportError:
    # OpenMV firmware ships without the on-device `mip` installer, so this script
    # can't run there. Point the user at the mip-free install paths instead of
    # dying with a bare "no module named 'mip'".
    print(
        "\n[tendrl] install_script.py needs the on-device 'mip' installer, but this "
        "firmware doesn't have it (OpenMV firmware ships without 'mip').\n"
        "         Install the SDK a mip-free way instead:\n"
        "           - Provision from the Contact Device tab or tendrl-dev-mcp "
        "(the SDK is pushed over serial), or\n"
        "           - run mip on your computer, which copies the files over USB:\n"
        "               mpremote mip install --target=/flash/lib "
        "https://app.tendrl.com/api/public/sdk/v1/latest/mpy/package.json\n"
    )
    raise SystemExit

CONFIG_FILE = "/config.json"
LIB_DIR = "/lib"
LIBRARY_CONFIG_FILE = "/lib/tendrl/config.json"
# Where the SDK is fetched from. Contact serves the build its own backend was
# deployed with, so a device gets a coherent SDK instead of whatever the repo's
# default branch happens to hold. Overridden per-environment by `app_url` in
# /config.json, so dev, dev-prod and self-hosted all work without editing this.
META = None  # populated from the backend in main()
DEFAULT_APP_URL = "https://app.tendrl.com"
SDK_PATH = "/api/public/sdk/v1/latest/mpy"
MAX_WIFI_RETRIES = 3
MAX_INSTALL_RETRIES = 3
WIFI_RETRY_DELAY = 5
INSTALL_RETRY_DELAY = 10

INSTALL_DB = True  # Set to False for minimal installation (no database)
INSTALL_FILES = False  # Set to True to add HTTP file transfer (minimal installs only; included in full)
INSTALL_STREAMING = False  # Set to True to include JPEG streaming support (adds ~25KB flash)
INSTALL_STREAMING_LITE = False  # Lighter streaming add-on (no permission pre-check)
INSTALL_VISION = False  # Set to True to include on-device vision (OpenMV only: csi/sensor, ml)

if INSTALL_DB:
    print("🗃️ Installing full Tendrl SDK with MicroTetherDB")
else:
    print("🗃️ Installing minimal Tendrl SDK (no database)")

if INSTALL_STREAMING:
    print("📹 Including JPEG streaming support")

if INSTALL_VISION:
    print("👁️ Including on-device vision support")

INCLUDE_DB = INSTALL_DB
INCLUDE_FILES = INSTALL_FILES or INSTALL_DB
INCLUDE_STREAMING = INSTALL_STREAMING
INCLUDE_STREAMING_LITE = INSTALL_STREAMING_LITE and not INSTALL_STREAMING
INCLUDE_VISION = INSTALL_VISION

def create_user_config_template():
    if not file_exists(CONFIG_FILE):
        template = {
            "api_key": "",
            "wifi_ssid": "",
            "wifi_pw": ""
        }
        try:
            with open(CONFIG_FILE, 'w') as f:
                formatted_json = "{\n \"api_key\": \"\",\n \"wifi_ssid\": \"\",\n \"wifi_pw\": \"\"\n}"
                f.write(formatted_json)
            print("✅ Created user config template at /config.json")
            print("⚠️ Please edit /config.json with your API key and WiFi credentials")
            print("📋 Note: API key ID and subject will be cached automatically after first connection")
        except Exception as e:
            print(f"❌ Failed to create user config template: {e}")
            return False
    return True

def load_config():
    try:
        with open(CONFIG_FILE) as f:
            return json.load(f)
    except Exception as e:
        print(f"❌ Failed to load config file: {e}")
        raise RuntimeError(f"Failed to load config file: {e}")

def connect_wifi(ssid, password, timeout=10):
    wlan = None
    try:
        wlan = network.WLAN(network.STA_IF)
        wlan.active(False)
        time.sleep(1)
        wlan.active(True)
        time.sleep(1)

        if wlan.isconnected():
            wlan.disconnect()
            time.sleep(1)

        for attempt in range(MAX_WIFI_RETRIES):
            try:
                print(f"🌐 Attempt {attempt + 1}/{MAX_WIFI_RETRIES}: Connecting to Wi-Fi: {ssid} ...")
                wlan.connect(ssid, password)

                start = time.time()
                while not wlan.isconnected():
                    if time.time() - start > timeout:
                        raise RuntimeError("Wi-Fi connection timed out")
                    time.sleep(0.5)

                if wlan.isconnected():
                    print("✅ Connected:", wlan.ifconfig())
                    return True
                else:
                    raise RuntimeError("Connection failed")

            except Exception as e:
                print(f"❌ Connection attempt {attempt + 1} failed: {e}")
                if attempt < MAX_WIFI_RETRIES - 1:
                    print(f"⏳ Waiting {WIFI_RETRY_DELAY} seconds before retrying...")
                    time.sleep(WIFI_RETRY_DELAY)
                else:
                    print("❌ All WiFi connection attempts failed")
                    return False

    except Exception as e:
        print(f"❌ WiFi setup failed: {e}")
        return False
    finally:
        if wlan and not wlan.isconnected():
            try:
                wlan.disconnect()
                wlan.active(False)
            except:
                pass

    return False

def file_exists(path):
    try:
        with open(path, 'r'):
            return True
    except:
        return False

def verify_installation():
    try:
        required_files = [
            "/lib/tendrl/__init__.py",
            "/lib/tendrl/client.py",
            "/lib/tendrl/config_manager.py",
            "/lib/tendrl/network_manager.py",
            "/lib/tendrl/queue_manager.py",
            "/lib/tendrl/mqtt_handler.py",
            "/lib/tendrl/lib/shutil.py",
            "/lib/tendrl/utils/__init__.py",
            "/lib/tendrl/utils/util_helpers.py",
            "/lib/tendrl/utils/http_helpers.py",
            "/lib/tendrl/manifest.py",
            "/lib/tendrl/config.json"
        ]

        if INCLUDE_DB:
            db_files = [
                "/lib/tendrl/lib/microtetherdb/__init__.py",
                "/lib/tendrl/lib/microtetherdb/db.py",
                "/lib/tendrl/lib/microtetherdb/core/__init__.py",
                "/lib/tendrl/lib/microtetherdb/core/exceptions.py",
                "/lib/tendrl/lib/microtetherdb/core/flush_manager.py",
                "/lib/tendrl/lib/microtetherdb/core/future.py",
                "/lib/tendrl/lib/microtetherdb/core/key_generator.py",
                "/lib/tendrl/lib/microtetherdb/core/query_engine.py",
                "/lib/tendrl/lib/microtetherdb/core/ttl_manager.py",
                "/lib/tendrl/lib/microtetherdb/core/ttl_expiry.py",
                "/lib/tendrl/lib/microtetherdb/core/utils.py"
            ]
            required_files.extend(db_files)
            required_files.extend([
                "/lib/tendrl/file_transfer.py",
                "/lib/tendrl/lib/zip_store.py",
            ])

        if INCLUDE_FILES and not INCLUDE_DB:
            required_files.extend([
                "/lib/tendrl/file_transfer.py",
                "/lib/tendrl/lib/zip_store.py",
            ])

        if INCLUDE_STREAMING:
            required_files.append("/lib/tendrl/streaming.py")
        elif INCLUDE_STREAMING_LITE:
            required_files.append("/lib/tendrl/streaming_lite.py")

        # The install may be source or precompiled depending on what the backend
        # serves, so accept either extension. Checking only .py would fail every
        # bytecode install here and burn all three retries on a good install.
        for file in required_files:
            if file_exists(file):
                continue
            if file.endswith(".py") and file_exists(file[:-3] + ".mpy"):
                continue
            print(f"❌ Required file not found: {file} (or its .mpy)")
            return False

        return True
    except Exception as e:
        print(f"❌ Error verifying installation: {e}")
        return False

def ensure_directory_exists(path):
    try:
        os.mkdir(path)
    except:
        pass

def create_library_config():
    try:
        if file_exists(LIBRARY_CONFIG_FILE):
            print("✅ Library config already exists at /lib/tendrl/config.json")
            return True

        ensure_directory_exists("/lib/tendrl")

        formatted_json = "{\n  \"tendrl_version\": \"0.1.0\",\n  \"app_url\": \"https://app.tendrl.com\",\n  \"api_key\": \"\",\n  \"wifi_ssid\": \"\",\n  \"wifi_pw\": \"\",\n  \"reset\": false,\n  \"mqtt_host\": \"mqtt.tendrl.com\",\n  \"mqtt_port\": 1883,\n  \"mqtt_ssl\": false\n}"

        with open(LIBRARY_CONFIG_FILE, 'w') as f:
            f.write(formatted_json)
        print("✅ Created library config at /lib/tendrl/config.json")
        return True
    except Exception as e:
        print(f"❌ Failed to create library config: {e}")
        return False

def ensure_required_directories():
    try:
        ensure_directory_exists("/lib")
        ensure_directory_exists("/lib/tendrl")
        ensure_directory_exists("/lib/tendrl/lib")

        if INCLUDE_DB:
            ensure_directory_exists("/lib/tendrl/lib/microtetherdb")
            ensure_directory_exists("/lib/tendrl/lib/microtetherdb/core")

        print("✅ Verified required directories")
        return True
    except Exception as e:
        print(f"❌ Failed to create required directories: {e}")
        return False

def sdk_base():
    """Where to fetch the SDK from, honouring this device's configured backend."""
    app_url = DEFAULT_APP_URL
    try:
        with open(CONFIG_FILE) as f:
            app_url = json.load(f).get("app_url") or DEFAULT_APP_URL
    except Exception:
        pass
    return app_url.rstrip("/") + SDK_PATH


def fetch_meta():
    """Ask the backend what it serves. None if unreachable — callers degrade."""
    try:
        import requests
    except ImportError:
        try:
            import urequests as requests
        except ImportError:
            return None
    try:
        r = requests.get(sdk_base().rsplit("/", 2)[0] + "/meta")
        if r.status_code != 200:
            print("⚠️ SDK metadata returned HTTP %s" % r.status_code)
            return None
        return r.json()
    except Exception as e:
        print("⚠️ Could not reach the SDK endpoint: %s" % e)
        return None


def check_bytecode_compatible(meta):
    """The served .mpy must match this firmware's bytecode version.

    A mismatch installs files the board physically cannot import, so fail loudly
    here rather than leave a device that looks provisioned and isn't.
    """
    if not meta:
        return True  # couldn't ask; let the install proceed and fail visibly
    served = meta.get("mpyVersion") or 0
    try:
        import sys
        ours = getattr(sys.implementation, "_mpy", 0) & 0xFF
    except Exception:
        ours = 0
    if not served or not ours or served == ours:
        return True
    print(
        "\n❌ This firmware speaks .mpy bytecode v%d but the server serves v%d.\n"
        "   Update MicroPython, or provision over serial from the Contact Device tab.\n"
        % (ours, served)
    )
    return False


def prune_managed(meta):
    """Delete previously-installed SDK files before writing the new ones.

    Mandatory, not tidy-up: MicroPython resolves `foo.py` before `foo.mpy`
    (py/builtinimport.c, stat_file_py_or_mpy), and mip never deletes. Without
    this, a board that already has a .py install keeps executing the stale source
    after a "successful" .mpy install, silently and indefinitely.

    Bounded by the server's managed set, so on-device databases and anything else
    the user put under lib/tendrl survive.
    """
    if not meta:
        print("⚠️ Skipping prune — no SDK metadata. A previous .py install may shadow the new files.")
        return
    managed = meta.get("managed") or []
    removed = 0
    dirs = []
    for rel in managed:
        full = LIB_DIR + "/" + rel
        try:
            os.remove(full)
            removed += 1
        except Exception:
            pass
        d = full.rsplit("/", 1)[0]
        while len(d) > len(LIB_DIR):
            if d not in dirs:
                dirs.append(d)
            d = d.rsplit("/", 1)[0]
    # Deepest first, and only when empty — lib/tendrl also holds device databases.
    dirs.sort(key=len)
    dirs.reverse()
    for d in dirs:
        try:
            if not os.listdir(d):
                os.rmdir(d)
        except Exception:
            pass
    if removed:
        print("🧹 Removed %d previously installed SDK file(s)" % removed)


def ensure_umqtt(base):
    """Install umqtt when the firmware doesn't already provide an MQTT client.

    tendrl/mqtt_handler.py wants umqtt.simple/umqtt.robust and falls back to
    OpenMV's built-in `mqtt`. The ESP32 port freezes umqtt into every build
    (ports/esp32/boards/manifest.py requires it); the rp2 port does not, so a
    stock Pico W has neither and the client comes up with MQTT silently
    disabled. umqtt is pure Python, so unlike btree it can just be installed.

    Gated on CAPABILITY, not board identity. A `platform == 'rp2'` test would
    both miss other ports that don't freeze umqtt and wrongly fire on custom
    Pico W firmware that does. ESP32 and OpenMV therefore pay nothing — and
    even if this did run there, sys.path is ['', '.frozen', '/lib'], so the
    frozen copy still wins.

    Installed from our own vendored copy (package-umqtt.json), not from
    micropython-lib: the SDK moved off third-party install sources precisely so
    that what a board runs matches the backend it talks to. It also means one
    reachable host instead of two, which matters on locked-down networks.
    """
    try:
        import umqtt.simple  # noqa: F401
        import umqtt.robust  # noqa: F401
        return True
    except ImportError:
        pass
    try:
        import mqtt  # noqa: F401  (OpenMV built-in)
        return True
    except ImportError:
        pass

    print("📨 No MQTT client in this firmware — installing umqtt...")
    try:
        mip.install(base + "/package-umqtt.json", target=LIB_DIR)
        print("   ✅ umqtt installed")
        return True
    except Exception as e:
        # Not fatal: everything except publishing still works, and the client
        # already degrades with a warning rather than crashing.
        print("   ⚠️ umqtt install failed (%s) — MQTT will be disabled" % e)
        return False


def install_tendrl():
    for attempt in range(MAX_INSTALL_RETRIES):
        try:
            if INCLUDE_DB:
                print(f"⬇️ Attempt {attempt + 1}/{MAX_INSTALL_RETRIES}: Installing full Tendrl SDK...")
            else:
                print(f"⬇️ Attempt {attempt + 1}/{MAX_INSTALL_RETRIES}: Installing minimal Tendrl SDK...")

            if not ensure_required_directories():
                raise RuntimeError("Failed to create required directories")

            base = sdk_base()
            print("📡 Source: %s" % base)
            # Clear any previous install first — see prune_managed(): a leftover
            # .py shadows the .mpy we are about to write.
            prune_managed(META)

            ensure_umqtt(base)

            if INCLUDE_DB:
                try:
                    import btree  # noqa: F401
                except ImportError:
                    # btree is a C module compiled into the firmware — it can't
                    # be installed. The rp2 port (Pico W) leaves it off, so the
                    # full tier would ship MicroTetherDB that can never import.
                    print("⚠️ This firmware has no 'btree' module, so MicroTetherDB "
                          "cannot run here.")
                    print("   Installing anyway, but pass client_db=False and expect "
                          "no persistent offline storage.")
                    print("   The minimal tier (INCLUDE_DB = False) is the better fit.")

            if INCLUDE_DB:
                mip.install(base + "/package.json", target=LIB_DIR)
            else:
                mip.install(base + "/package-minimal.json", target=LIB_DIR)

            if INCLUDE_FILES and not INCLUDE_DB:
                print("📁 Installing HTTP file transfer module...")
                mip.install(base + "/package-files.json", target=LIB_DIR)

            # Install streaming module if requested
            if INCLUDE_STREAMING:
                print("📹 Installing JPEG streaming module...")
                mip.install(base + "/package-streaming.json", target=LIB_DIR)
            elif INCLUDE_STREAMING_LITE:
                print("📹 Installing JPEG streaming module (lite)...")
                mip.install(base + "/package-streaming-lite.json", target=LIB_DIR)

            # Install on-device vision module if requested (OpenMV only)
            if INCLUDE_VISION:
                print("👁️ Installing on-device vision module...")
                mip.install(base + "/package-vision.json", target=LIB_DIR)

            if not create_library_config():
                raise RuntimeError("Failed to create library config")

            if verify_installation():
                if INCLUDE_DB:
                    print("✅ Full Tendrl SDK installed and verified successfully")
                    print("📊 Includes MicroTetherDB for local data storage")
                else:
                    print("✅ Minimal Tendrl SDK installed and verified successfully")
                    print("⚠️ Note: Local database features disabled (client_db=False required)")
                return True
            else:
                raise RuntimeError("Installation verification failed")

        except Exception as e:
            print(f"❌ Installation attempt {attempt + 1} failed: {e}")
            if attempt < MAX_INSTALL_RETRIES - 1:
                print(f"⏳ Waiting {INSTALL_RETRY_DELAY} seconds before retrying...")
                time.sleep(INSTALL_RETRY_DELAY)
            else:
                print("❌ All installation attempts failed")
                return False
    return False

def main():
    try:
        if not file_exists(CONFIG_FILE):
            print("⚠️ User config not found")
            if not create_user_config_template():
                print("❌ Failed to create user config template")
                return
            print("Required fields:")
            print("  - wifi_ssid: Your WiFi network name")
            print("  - wifi_pw: Your WiFi password")
            print("\nAfter filling in these details, run this script again.")
            return

        config = load_config()

        ssid = config.get("wifi_ssid")
        pw = config.get("wifi_pw")
        if not ssid or not pw:
            print("⚠️ Missing Wi-Fi credentials in config.json")
            print("Please edit /config.json with your WiFi credentials and API key")
            return

        if not ensure_required_directories():
            print("❌ Failed to create required directories")
            return

        if not connect_wifi(ssid, pw):
            print("❌ Failed to establish WiFi connection after all retries")
            return

        # Ask the backend what it serves: the bytecode version we must match, and
        # the file set to clear before installing. Needs the network, so it comes
        # after the Wi-Fi join.
        global META
        META = fetch_meta()
        if META:
            print("📦 SDK %s from this backend" % META.get("version", "?"))
        if not check_bytecode_compatible(META):
            return

        if not install_tendrl():
            print("❌ Failed to install tendrl after all retries")
            return

        print("✨ Installation completed successfully!")
        if INCLUDE_DB:
            print("📊 MicroTetherDB is available for local data storage")
            print("💡 Use client_db=True in Client() constructor (default)")
        else:
            print("⚠️ MicroTetherDB not installed - use client_db=False in Client() constructor")
            print("💡 This saves ~35KB flash space but disables local database features")
        if INCLUDE_STREAMING:
            print("📹 JPEG streaming module installed")
            print("💡 Use client.start_streaming() for camera streaming")
        elif INCLUDE_STREAMING_LITE:
            print("📹 JPEG streaming module (lite) installed")
            print("💡 Use client.start_streaming() for camera streaming")
        else:
            print("💡 To enable JPEG streaming, set INSTALL_STREAMING=True in install_script.py")
            print("   Streaming adds ~25KB flash storage and works with both minimal and full installations")
        if INCLUDE_VISION:
            print("👁️ On-device vision module installed")
            print("💡 Use tendrl.vision: Camera, Watch, and the detectors (roi_brightness, person, ...)")
        else:
            print("💡 To enable on-device vision, set INSTALL_VISION=True in install_script.py")
            print("   OpenMV only — needs the csi/sensor module (and ml for model detectors)")
        if not INCLUDE_DB and not INCLUDE_FILES:
            print("💡 To enable file transfer, set INSTALL_FILES=True in install_script.py")
        print("⚠️ If you haven't already, please edit /config.json with your API key")
    except Exception as e:
        print(f"❌ Unexpected error: {e}")
        return

if __name__ == "__main__":
    main()