import pyray as rl
import qrcode
import numpy as np
import time
import jwt
import os
from datetime import datetime, timedelta, UTC

from openpilot.common.api.base import BaseApi
from openpilot.common.swaglog import cloudlog
from openpilot.common.params import Params
from openpilot.iqpilot.konn3kt.registration import get_or_create_dongle_id, ensure_dev_pairing_identity
from openpilot.system.hardware import HARDWARE, PC
from openpilot.system.ui.widgets import Widget
from openpilot.system.ui.lib.application import FontWeight, gui_app
from openpilot.system.ui.lib.multilang import tr
from openpilot.system.ui.lib.wrap_text import wrap_text
from openpilot.system.ui.lib.text_measure import measure_text_cached
from openpilot.system.ui.widgets.button import IconButton
from openpilot.selfdrive.ui.ui_state import ui_state


class PairingDialog(Widget):
  """Dialog for device pairing with QR code."""

  QR_REFRESH_INTERVAL = 300  # 5 minutes in seconds

  def __init__(self):
    super().__init__()
    self.params = Params()
    self.qr_texture: rl.Texture | None = None
    self.last_qr_generation = float('-inf')
    self._close_btn = IconButton(gui_app.texture("icons/iq/close.png", 80, 80))
    self._close_btn.set_click_callback(lambda: gui_app.set_modal_overlay(None))

  def _get_pairing_url(self) -> str:
    dev_pairing = PC and os.getenv("KONN3KT_DEV_PAIRING") == "1"
    if dev_pairing:
      try:
        ensure_dev_pairing_identity(self.params, force_reset=os.getenv("KONN3KT_DEV_PAIRING_RESET") == "1")
      except Exception:
        return "error://dev_identity_setup_failed"

    try:
      imei1 = HARDWARE.get_imei(0) or ""
    except Exception as e:
      cloudlog.warning(f"Failed to get imei1: {e}")
      imei1 = ""

    try:
      imei2 = HARDWARE.get_imei(1) or ""
    except Exception as e:
      cloudlog.warning(f"Failed to get imei2: {e}")
      imei2 = ""

    try:
      algorithm, private_key, public_key = BaseApi.get_key_pair()
      if not private_key or not algorithm:
        cloudlog.error("No device keys found")
        return "error://keys_not_found"

      dongle_id = get_or_create_dongle_id(self.params, prefer_readonly=True)

      try:
        serial = HARDWARE.get_serial() or ""
      except Exception as e:
        cloudlog.warning(f"Failed to get serial: {e}")
        serial = ""
      if not serial:
        serial = (self.params.get("HardwareSerial") or "") if dev_pairing else ""
      if not serial:
        cloudlog.error("No hardware serial found, cannot generate pairing token")
        return "error://serial_not_found"

      now = datetime.now(UTC).replace(tzinfo=None)
      payload = {
        'identity': dongle_id,
        'nbf': now,
        'iat': now,
        'imei': imei1,
        'imei2': imei2,
        'serial': serial,
        'public_key': public_key,
        'register': True,
        'exp': now + timedelta(hours=1),
      }

      try:
        token = jwt.encode(payload, private_key, algorithm=algorithm)
      except Exception as e:
        cloudlog.warning(f"jwt.encode failed ({e}), retrying with normalized key")
        try:
          from cryptography.hazmat.primitives import serialization
          key_bytes = private_key.encode("utf-8") if isinstance(private_key, str) else private_key
          try:
            key_obj = serialization.load_pem_private_key(key_bytes, password=None)
          except Exception:
            key_obj = serialization.load_ssh_private_key(key_bytes, password=None)
          token = jwt.encode(payload, key_obj, algorithm=algorithm)
        except Exception as e2:
          cloudlog.error(f"Failed to generate pairing token: {e2}")
          return "error://token_generation_failed"
      if isinstance(token, bytes):
        token = token.decode('utf8')
      return f"https://konn3kt.com/?pair={token}"
    except FileNotFoundError as e:
      cloudlog.error(f"Key files not found: {e}")
      return "error://keys_not_found"
    except Exception as e:
      cloudlog.error(f"Failed to generate pairing token: {e}")
      return "error://token_generation_failed"

  def _generate_qr_code(self) -> None:
    try:
      url = self._get_pairing_url()
      if url.startswith("error://"):
        cloudlog.warning(f"Cannot generate QR code: {url}")
        self.qr_texture = None
        return

      qr = qrcode.QRCode(version=1, error_correction=qrcode.constants.ERROR_CORRECT_L, box_size=10, border=4)
      qr.add_data(url)
      qr.make(fit=True)

      pil_img = qr.make_image(fill_color="black", back_color="white").convert('RGBA')
      img_array = np.array(pil_img, dtype=np.uint8)

      if self.qr_texture and self.qr_texture.id != 0:
        rl.unload_texture(self.qr_texture)

      rl_image = rl.Image()
      rl_image.data = rl.ffi.cast("void *", img_array.ctypes.data)
      rl_image.width = pil_img.width
      rl_image.height = pil_img.height
      rl_image.mipmaps = 1
      rl_image.format = rl.PixelFormat.PIXELFORMAT_UNCOMPRESSED_R8G8B8A8

      self.qr_texture = rl.load_texture_from_image(rl_image)
    except Exception:
      cloudlog.exception("QR code generation failed")
      self.qr_texture = None

  def _check_qr_refresh(self) -> None:
    current_time = time.monotonic()
    if current_time - self.last_qr_generation >= self.QR_REFRESH_INTERVAL:
      self._generate_qr_code()
      self.last_qr_generation = current_time

  def _update_state(self):
    if ui_state.prime_state.is_paired():
      gui_app.set_modal_overlay(None)

  def _render(self, rect: rl.Rectangle) -> int:
    rl.clear_background(rl.Color(20, 21, 24, 255))

    self._check_qr_refresh()

    margin = 70
    content_rect = rl.Rectangle(rect.x + margin, rect.y + margin, rect.width - 2 * margin, rect.height - 2 * margin)
    y = content_rect.y

    # Close button
    close_size = 80
    pad = 20
    close_rect = rl.Rectangle(content_rect.x - pad, y - pad, close_size + pad * 2, close_size + pad * 2)
    self._close_btn.render(close_rect)

    y += close_size + 40

    # Title
    title = tr("Pair your device to your Konn3kt account")
    title_font = gui_app.font(FontWeight.NORMAL)
    left_width = int(content_rect.width * 0.5 - 15)

    title_wrapped = wrap_text(title_font, title, 75, left_width)
    rl.draw_text_ex(title_font, "\n".join(title_wrapped), rl.Vector2(content_rect.x, y), 75, 0.0, rl.WHITE)
    y += len(title_wrapped) * 75 + 60

    # Two columns: instructions and QR code
    remaining_height = content_rect.height - (y - content_rect.y)
    right_width = content_rect.width // 2 - 20

    # Instructions
    self._render_instructions(rl.Rectangle(content_rect.x, y, left_width, remaining_height))

    # QR code
    qr_size = min(right_width, content_rect.height) - 40
    qr_x = content_rect.x + left_width + 40 + (right_width - qr_size) // 2
    qr_y = content_rect.y
    self._render_qr_code(rl.Rectangle(qr_x, qr_y, qr_size, qr_size))

    return -1

  def _render_instructions(self, rect: rl.Rectangle) -> None:
    instructions = [
      tr("Open the Konn3kt app on your phone"),
      tr("Tap \"Add Device\" and scan the QR code on the right"),
      tr("Follow the prompts in the app to finish pairing"),
    ]

    font = gui_app.font(FontWeight.BOLD)
    y = rect.y

    for i, text in enumerate(instructions):
      circle_radius = 25
      circle_x = rect.x + circle_radius + 15
      text_x = rect.x + circle_radius * 2 + 40
      text_width = rect.width - (circle_radius * 2 + 40)

      wrapped = wrap_text(font, text, 47, int(text_width))
      text_height = len(wrapped) * 47
      circle_y = y + text_height // 2

      # Circle and number
      rl.draw_circle(int(circle_x), int(circle_y), circle_radius, rl.Color(16, 185, 169, 255))
      number = str(i + 1)
      number_size = measure_text_cached(font, number, 30)
      rl.draw_text_ex(font, number, (int(circle_x - number_size.x // 2), int(circle_y - number_size.y // 2)), 30, 0, rl.WHITE)

      # Text
      rl.draw_text_ex(font, "\n".join(wrapped), rl.Vector2(text_x, y), 47, 0.0, rl.WHITE)
      y += text_height + 50

  def _render_qr_code(self, rect: rl.Rectangle) -> None:
    # White card: QR codes must stay light to scan, and it reads as an intentional panel on the dark theme.
    rl.draw_rectangle_rounded(rect, 0.06, 20, rl.Color(245, 245, 245, 255))

    if not self.qr_texture:
      error_font = gui_app.font(FontWeight.BOLD)
      msg = tr("QR Code Error")
      ms = measure_text_cached(error_font, msg, 34)
      pos = rl.Vector2(rect.x + (rect.width - ms.x) / 2, rect.y + (rect.height - ms.y) / 2)
      rl.draw_text_ex(error_font, msg, pos, 34, 0.0, rl.Color(200, 60, 52, 255))
      return

    pad = 28
    inner = rl.Rectangle(rect.x + pad, rect.y + pad, rect.width - 2 * pad, rect.height - 2 * pad)
    source = rl.Rectangle(0, 0, self.qr_texture.width, self.qr_texture.height)
    rl.draw_texture_pro(self.qr_texture, source, inner, rl.Vector2(0, 0), 0, rl.WHITE)

  def __del__(self):
    if self.qr_texture and self.qr_texture.id != 0:
      rl.unload_texture(self.qr_texture)


if __name__ == "__main__":
  gui_app.init_window("pairing device")
  pairing = PairingDialog()
  try:
    for _ in gui_app.render():
      result = pairing.render(rl.Rectangle(0, 0, gui_app.width, gui_app.height))
      if result != -1:
        break
  finally:
    del pairing