#!/usr/bin/env python3
"""
Copyright (c) IQ.Lvbs, apart of Project Teal Lvbs, All Rights Reserved, licensed under https://konn3kt.com/tos
"""
import asyncio
import hashlib
import os
import time
from pathlib import Path

import aiohttp
from cereal import custom
from openpilot.common.realtime import Ratekeeper
from openpilot.common.time_helpers import system_time_valid
from openpilot.iqpilot._proprietary_loader import ProprietaryModuleMissing, load_private_module
from openpilot.common.swaglog import cloudlog
from openpilot.system.hardware.hw import Paths

_TIME_SYNC_WAIT_TIMEOUT_S = 30.0
_TIME_SYNC_POLL_S = 0.5


def _wait_for_valid_clock(timeout: float = _TIME_SYNC_WAIT_TIMEOUT_S) -> None:
  if system_time_valid():
    return
  cloudlog.warning("models_manager: system clock not yet valid, waiting for NTP before fetching")
  deadline = time.monotonic() + timeout
  while time.monotonic() < deadline:
    if system_time_valid():
      cloudlog.warning("models_manager: system clock is now valid, resuming")
      return
    time.sleep(_TIME_SYNC_POLL_S)
  cloudlog.warning("models_manager: gave up waiting for a valid clock, proceeding anyway")

try:
  load_private_module(__name__, "iqpilot_private.models.manager")
  _BaseIQModelManager = IQModelManager  # noqa: F821
except ProprietaryModuleMissing:
  from iqpilot.models_private_src.manager import IQModelManager as _BaseIQModelManager

from openpilot.iqpilot.selfdrive.iqmodeld.models.git_auth import get_aiohttp_auth
from openpilot.iqpilot.selfdrive.iqmodeld.models.helpers import (
  bundle_files_ready,
  get_active_bundle,
  get_runtime_bundle_upgrade,
  is_default_bundle,
  persist_active_bundle,
)


_ACTIVE_BUNDLE_KEY = "ModelManager_ActiveBundle"
_DOWNLOAD_INDEX_KEY = "ModelManager_DownloadIndex"
_RUNNER_CACHE_KEY = "ModelRunnerTypeCache"


class IQModelManager(_BaseIQModelManager):
  def __init__(self):
    super().__init__()
    self._validated_active_key: tuple[tuple[str, str], ...] | None = None
    self._manifest_refresh_key: tuple[tuple[str, str], ...] | None = None

  @staticmethod
  def _bundle_index(bundle) -> int | None:
    try:
      return int(getattr(bundle, "index", -1))
    except (TypeError, ValueError):
      return None

  @staticmethod
  def _bundle_files(bundle) -> list[tuple[str, str]]:
    files = []
    for model in getattr(bundle, "models", []) or []:
      for artifact in (getattr(model, "metadata", None), getattr(model, "artifact", None)):
        filename = getattr(artifact, "fileName", "") if artifact is not None else ""
        if not filename:
          continue
        download_uri = getattr(artifact, "downloadUri", None)
        sha256 = getattr(download_uri, "sha256", "") if download_uri is not None else ""
        files.append((filename, sha256 or ""))
    return files

  @staticmethod
  def _safe_model_path(filename: str) -> Path | None:
    if not filename or os.path.basename(filename) != filename:
      cloudlog.warning(f"Ignoring unsafe model filename {filename!r}")
      return None

    root = Path(Paths.model_root()).resolve()
    path = (root / filename).resolve()
    try:
      path.relative_to(root)
    except ValueError:
      cloudlog.warning(f"Ignoring model path outside model root {path}")
      return None
    return path

  @staticmethod
  def _verify_file_sync(path: Path, expected_hash: str) -> bool:
    if not path.is_file():
      return False
    if not expected_hash:
      return True

    sha256_hash = hashlib.sha256()
    with open(path, "rb") as f:
      for chunk in iter(lambda: f.read(1024 * 1024), b""):
        sha256_hash.update(chunk)
    return sha256_hash.hexdigest().lower() == expected_hash.lower()

  def _bundle_validation_key(self, bundle) -> tuple[tuple[str, str], ...]:
    return tuple(self._bundle_files(bundle))

  def _bundle_files_valid(self, bundle) -> bool:
    for filename, expected_hash in self._bundle_files(bundle):
      path = self._safe_model_path(filename)
      if path is None or not self._verify_file_sync(path, expected_hash):
        return False
    return True

  def _remove_bundle_files(self, bundle) -> None:
    for filename, _expected_hash in self._bundle_files(bundle):
      path = self._safe_model_path(filename)
      if path is None:
        continue
      for candidate in (path, Path(f"{path}.download")):
        try:
          if candidate.is_file():
            candidate.unlink()
        except OSError as e:
          cloudlog.exception(f"Failed to remove model artifact {candidate}: {e}")

  def _find_available_bundle(self, target):
    target_index = self._bundle_index(target)
    target_ref = getattr(target, "ref", None)
    target_internal = getattr(target, "internalName", None)
    target_display = getattr(target, "displayName", None)

    for bundle in self.available_models:
      if target_index is not None and self._bundle_index(bundle) == target_index:
        return bundle
      if target_ref and getattr(bundle, "ref", None) == target_ref:
        return bundle
      if target_internal and getattr(bundle, "internalName", None) == target_internal:
        return bundle
      if target_display and getattr(bundle, "displayName", None) == target_display:
        return bundle
    return None

  def _bundle_matches(self, left, right) -> bool:
    if left is None or right is None:
      return False

    left_index = self._bundle_index(left)
    right_index = self._bundle_index(right)
    if left_index is not None and right_index is not None and left_index == right_index:
      return True

    for attr in ("ref", "internalName", "displayName"):
      left_value = getattr(left, attr, None)
      if left_value and left_value == getattr(right, attr, None):
        return True
    return False

  def _clear_active_bundle(self) -> None:
    self.params.remove(_ACTIVE_BUNDLE_KEY)
    self.params.remove(_RUNNER_CACHE_KEY)
    self.active_bundle = None
    self._validated_active_key = None

  def _download_request_matches(self, bundle) -> bool:
    bundle_index = self._bundle_index(bundle)
    return bundle_index is not None and self._download_index() == bundle_index

  def _queue_active_redownload_if_invalid(self) -> None:
    if self.active_bundle is None:
      self._validated_active_key = None
      return

    validation_key = self._bundle_validation_key(self.active_bundle)
    if validation_key == self._validated_active_key:
      return

    if self._bundle_files_valid(self.active_bundle):
      self._validated_active_key = validation_key
      return

    bundle = self._find_available_bundle(self.active_bundle) or self.active_bundle
    bundle_index = self._bundle_index(bundle)
    cloudlog.warning(f"Active model {_display_bundle_name(self.active_bundle)} is missing or corrupt; queueing redownload")
    self._remove_bundle_files(bundle)
    self._clear_active_bundle()
    if bundle_index is not None and self._download_index() is None:
      self.params.put(_DOWNLOAD_INDEX_KEY, bundle_index)

  def _find_manifest_counterpart(self, target):
    # never match by index: indexes shift between manifest generations, and a
    # positional match could redownload a different model than the user selected
    for attr in ("ref", "internalName", "displayName"):
      value = getattr(target, attr, None)
      if not value:
        continue
      for bundle in self.available_models:
        if getattr(bundle, attr, None) == value:
          return bundle
    return None

  def _queue_active_manifest_refresh(self) -> None:
    active = self.active_bundle
    if active is None or is_default_bundle(active):
      return
    if self._download_index() is not None:
      return

    counterpart = self._find_manifest_counterpart(active)
    if counterpart is None:
      return
    counterpart_index = self._bundle_index(counterpart)
    if counterpart_index is None:
      return

    active_files = dict(self._bundle_files(active))
    stale = False
    for filename, sha in self._bundle_files(counterpart):
      if not sha:
        continue
      active_sha = active_files.get(filename)
      # an empty recorded hash can't prove a mismatch, so it never triggers a redownload
      if active_sha is None or (active_sha and active_sha.lower() != sha.lower()):
        stale = True
        break
    if not stale:
      self._manifest_refresh_key = None
      return

    # the manifest may be an expired offline cache, so keep the active bundle and its
    # files in place: the download flow replaces artifacts atomically and only persists
    # the counterpart as active once everything landed. One attempt per bundle per run
    # so a dead network doesn't turn the 1Hz loop into a download-retry storm.
    key = self._bundle_validation_key(active)
    if key == self._manifest_refresh_key:
      return
    self._manifest_refresh_key = key

    cloudlog.warning(f"Active model {_display_bundle_name(active)} artifacts are stale vs current manifest; queueing redownload")
    self.params.put(_DOWNLOAD_INDEX_KEY, counterpart_index)

  async def _download_file(self, url: str, path: str, model) -> None:
    temp_path = f"{path}.download"
    self._download_start_times[model.fileName] = time.monotonic()

    try:
      if os.path.exists(temp_path):
        os.remove(temp_path)

      async with aiohttp.ClientSession(auth=get_aiohttp_auth()) as session:
        async with session.get(url) as response:
          response.raise_for_status()
          total_size = int(response.headers.get("content-length", 0))
          bytes_downloaded = 0

          with open(temp_path, "wb") as f:
            async for chunk in response.content.iter_chunked(self._chunk_size):
              f.write(chunk)
              bytes_downloaded += len(chunk)

              if self._download_index() is None:
                raise Exception("Download cancelled")

              if total_size > 0:
                progress = (bytes_downloaded / total_size) * 100
                model.downloadProgress.status = custom.IQModelManager.DownloadStatus.downloading
                model.downloadProgress.progress = progress
                model.downloadProgress.eta = self._calculate_eta(model.fileName, progress)
                self._report_status()

            f.flush()
            os.fsync(f.fileno())

      os.replace(temp_path, path)

    except Exception:
      if os.path.exists(temp_path):
        os.remove(temp_path)
      raise

    finally:
      self._download_start_times.pop(model.fileName, None)

  async def _download_bundle(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
    self.selected_bundle = model_bundle
    self.selected_bundle.status = custom.IQModelManager.DownloadStatus.downloading
    os.makedirs(destination_path, exist_ok=True)

    try:
      if not self._download_request_matches(model_bundle):
        raise RuntimeError("Download cancelled")

      tasks = [self._process_model(model, destination_path) for model in self.selected_bundle.models]
      await asyncio.gather(*tasks)

      if not self._download_request_matches(model_bundle):
        raise RuntimeError("Download cancelled")

      self.active_bundle = self.selected_bundle
      self.active_bundle.status = custom.IQModelManager.DownloadStatus.downloaded
      self.params.put(_ACTIVE_BUNDLE_KEY, self.active_bundle.to_dict())
      self.params.remove(_RUNNER_CACHE_KEY)
      self.selected_bundle = None

    except Exception:
      if self._download_request_matches(model_bundle) and self.selected_bundle is not None:
        self.selected_bundle.status = custom.IQModelManager.DownloadStatus.failed
      else:
        self.selected_bundle = None
      raise

    finally:
      self._report_status()

  def download(self, model_bundle: custom.IQModelManager.ModelBundle, destination_path: str) -> None:
    asyncio.run(self._download_bundle(model_bundle, destination_path))

  def _queue_tinygrad_upgrade(self) -> None:
    if self.active_bundle is None:
      return

    replacement = get_runtime_bundle_upgrade(self.active_bundle, self.params, self.available_models)
    if replacement is None or replacement is self.active_bundle:
      return

    if bundle_files_ready(replacement):
      persist_active_bundle(self.params, replacement)
      self.active_bundle = replacement
      return

    if self._download_index() is None and getattr(replacement, "index", None) is not None:
      self.params.put("ModelManager_DownloadIndex", int(replacement.index))
      cloudlog.warning(f"Queued tinygrad upgrade for retired bundle {getattr(self.active_bundle, 'internalName', '<unknown>')}")

  def main_thread(self) -> None:
    _wait_for_valid_clock()
    rk = Ratekeeper(1, print_delay_threshold=None)

    while True:
      try:
        # before NTP the TLS cert reads "not yet valid" and every fetch SSL-fails; one line, not spam
        if not system_time_valid():
          if not getattr(self, "_ntp_wait_logged", False):
            cloudlog.warning("models_manager: waiting for NTP before fetching (system clock not valid)")
            self._ntp_wait_logged = True
          rk.keep_time()
          continue
        self._ntp_wait_logged = False

        self.available_models = self.model_fetcher.get_available_bundles()
        self.active_bundle = get_active_bundle(self.params)
        self._queue_active_redownload_if_invalid()
        self._queue_tinygrad_upgrade()
        self._queue_active_manifest_refresh()

        if (index_to_download := self._download_index()) is not None:
          if model_to_download := next((model for model in self.available_models if model.index == index_to_download), None):
            try:
              self.download(model_to_download, Paths.model_root())
            except Exception as e:
              cloudlog.exception(e)
            finally:
              self.params.remove("ModelManager_DownloadIndex")
              self.selected_bundle = None

        if self.params.get("ModelManager_ClearCache"):
          self.clear_model_cache()
          self.params.remove("ModelManager_ClearCache")

        self._report_status()
        rk.keep_time()

      except Exception as e:
        cloudlog.exception(f"Error in main thread: {str(e)}")
        rk.keep_time()


def _display_bundle_name(bundle) -> str:
  return getattr(bundle, "internalName", None) or getattr(bundle, "displayName", None) or "<unknown>"


def main():
  IQModelManager().main_thread()

if __name__ == "__main__":
  main()