---
# JupyterLab Deployment with OAuth2-Proxy Sidecar
# Deploys JupyterLab as systemd service with Casdoor SSO via oauth2-proxy
# Red Panda Approved

- name: Deploy JupyterLab
  hosts: ubuntu
  become: true
  tasks:
    - name: Check if host has jupyterlab service
      ansible.builtin.set_fact:
        has_jupyterlab_service: "{{'jupyterlab' in services}}"

    - name: Skip hosts without jupyterlab service
      ansible.builtin.meta: end_host
      when: not has_jupyterlab_service

    # =========================================================================
    # System Dependencies
    # =========================================================================
    - name: Install system dependencies
      ansible.builtin.apt:
        name:
          - python3
          - python3-venv
          - python3-dev
          - python3-pip
          - nodejs
          - npm
          - graphviz
          - git
          - curl
        state: present
        update_cache: true

    # =========================================================================
    # User Setup
    # =========================================================================
    - name: Ensure jupyterlab user exists
      ansible.builtin.user:
        name: "{{ jupyterlab_user }}"
        group: "{{ jupyterlab_group }}"
        shell: /bin/bash
        create_home: true
        state: present

    - name: Create Notebooks directory
      ansible.builtin.file:
        path: "{{ jupyterlab_notebook_dir }}"
        owner: "{{ jupyterlab_user }}"
        group: "{{ jupyterlab_group }}"
        state: directory
        mode: '0755'

    - name: Create JupyterLab config directory
      ansible.builtin.file:
        path: /etc/jupyterlab
        owner: root
        group: "{{ jupyterlab_group }}"
        state: directory
        mode: '0755'

    - name: Create JupyterLab log directory
      ansible.builtin.file:
        path: /var/log/jupyterlab
        owner: "{{ jupyterlab_user }}"
        group: "{{ jupyterlab_group }}"
        state: directory
        mode: '0755'

    # =========================================================================
    # Python Virtual Environment
    # =========================================================================
    - name: Create virtual environment directory
      ansible.builtin.file:
        path: "{{ jupyterlab_venv_dir }}"
        owner: "{{ jupyterlab_user }}"
        group: "{{ jupyterlab_group }}"
        state: directory
        mode: '0755'

    - name: Create virtual environment for JupyterLab
      become_user: "{{ jupyterlab_user }}"
      ansible.builtin.command:
        cmd: "python3 -m venv {{ jupyterlab_venv_dir }}"
        creates: "{{ jupyterlab_venv_dir }}/bin/activate"

    - name: Upgrade pip in virtual environment
      become_user: "{{ jupyterlab_user }}"
      ansible.builtin.pip:
        name:
          - pip
          - wheel
          - setuptools
        state: latest
        virtualenv: "{{ jupyterlab_venv_dir }}"

    - name: Install JupyterLab and core packages
      become_user: "{{ jupyterlab_user }}"
      ansible.builtin.pip:
        name:
          - jupyterlab
          - jupyter-ai[all]
          - langchain-ollama
          - matplotlib
          - plotly
          - jupyter_contrib_nbextensions
          - "jsonschema[format-nongpl]"
          - python-mermaid
          - ipywidgets
        state: present
        virtualenv: "{{ jupyterlab_venv_dir }}"
      notify: restart jupyterlab

    # =========================================================================
    # Configuration Files
    # =========================================================================
    - name: Template JupyterLab configuration
      ansible.builtin.template:
        src: jupyter_lab_config.py.j2
        dest: /etc/jupyterlab/jupyter_lab_config.py
        owner: root
        group: "{{ jupyterlab_group }}"
        mode: '0644'
      notify: restart jupyterlab

    - name: Template JupyterLab systemd service
      ansible.builtin.template:
        src: jupyterlab.service.j2
        dest: /etc/systemd/system/jupyterlab.service
        owner: root
        group: root
        mode: '0644'
      notify:
        - reload systemd
        - restart jupyterlab

    # =========================================================================
    # OAuth2-Proxy Sidecar
    # =========================================================================
    - name: Create oauth2-proxy directory
      ansible.builtin.file:
        path: "{{ jupyterlab_oauth2_proxy_dir }}"
        owner: root
        group: root
        state: directory
        mode: '0755'

    - name: Download oauth2-proxy binary
      ansible.builtin.get_url:
        url: "https://github.com/oauth2-proxy/oauth2-proxy/releases/download/v{{ jupyterlab_oauth2_proxy_version }}/oauth2-proxy-v{{ jupyterlab_oauth2_proxy_version }}.linux-amd64.tar.gz"
        dest: "/tmp/oauth2-proxy-v{{ jupyterlab_oauth2_proxy_version }}.tar.gz"
        mode: '0644'

    - name: Extract oauth2-proxy binary
      ansible.builtin.unarchive:
        src: "/tmp/oauth2-proxy-v{{ jupyterlab_oauth2_proxy_version }}.tar.gz"
        dest: /tmp
        remote_src: true
        creates: "/tmp/oauth2-proxy-v{{ jupyterlab_oauth2_proxy_version }}.linux-amd64/oauth2-proxy"

    - name: Install oauth2-proxy binary
      ansible.builtin.copy:
        src: "/tmp/oauth2-proxy-v{{ jupyterlab_oauth2_proxy_version }}.linux-amd64/oauth2-proxy"
        dest: /usr/local/bin/oauth2-proxy
        owner: root
        group: root
        mode: '0755'
        remote_src: true

    - name: Template oauth2-proxy configuration
      ansible.builtin.template:
        src: oauth2-proxy-jupyter.cfg.j2
        dest: "{{ jupyterlab_oauth2_proxy_dir }}/oauth2-proxy.cfg"
        owner: root
        group: root
        mode: '0600'
      notify: restart oauth2-proxy-jupyter

    - name: Template oauth2-proxy systemd service
      ansible.builtin.template:
        src: oauth2-proxy-jupyter.service.j2
        dest: /etc/systemd/system/oauth2-proxy-jupyter.service
        owner: root
        group: root
        mode: '0644'
      notify:
        - reload systemd
        - restart oauth2-proxy-jupyter

    # =========================================================================
    # Service Management
    # =========================================================================
    - name: Enable and start JupyterLab service
      ansible.builtin.systemd:
        name: jupyterlab
        enabled: true
        state: started
        daemon_reload: true

    - name: Enable and start OAuth2-Proxy service
      ansible.builtin.systemd:
        name: oauth2-proxy-jupyter
        enabled: true
        state: started
        daemon_reload: true

  handlers:
    - name: reload systemd
      ansible.builtin.systemd:
        daemon_reload: true

    - name: restart jupyterlab
      ansible.builtin.systemd:
        name: jupyterlab
        state: restarted

    - name: restart oauth2-proxy-jupyter
      ansible.builtin.systemd:
        name: oauth2-proxy-jupyter
        state: restarted