#!/bin/bash
# setup-audio.sh — Set up barepi audio (TLV320DAC3100 + TDK T5848 mics) on a fresh Pi
# Usage: sudo ./setup-audio.sh
set -euo pipefail

if [ "$(id -u)" -ne 0 ]; then
    echo "Run as root: sudo ./setup-audio.sh"
    exit 1
fi

TARGET_USER="${TARGET_USER:-${SUDO_USER:-}}"
if [ -z "$TARGET_USER" ] || [ "$TARGET_USER" = "root" ]; then
    echo "Could not determine target user." >&2
    echo "Run with sudo from the target account or pass TARGET_USER=<user>." >&2
    exit 1
fi

TARGET_HOME="$(getent passwd "$TARGET_USER" | cut -d: -f6 || true)"
if [ -z "$TARGET_HOME" ] || [ ! -d "$TARGET_HOME" ]; then
    echo "Could not determine home directory for user: $TARGET_USER" >&2
    exit 1
fi

CONFIG="/boot/firmware/config.txt"
OVERLAYS="/boot/firmware/overlays"
RESTORE_HELPER="/usr/local/sbin/barepi-audio-restore"
RESTORE_SERVICE="/etc/systemd/system/barepi-audio-restore.service"

echo "=== barepi audio setup ==="
echo "User: $TARGET_USER  Home: $TARGET_HOME"
echo ""

# 1. Install packages
echo "[1/7] Installing packages..."
apt-get update
apt-get install -y alsa-utils libasound2-plugins swh-plugins i2c-tools device-tree-compiler
echo "      Done."

# 2. Compile and install device tree overlay
echo "[2/7] Installing device tree overlay..."
DTS=$(mktemp /tmp/i2s-audio-XXXX.dts)
trap 'rm -f "$DTS"' EXIT
cat > "$DTS" << 'DTS_EOF'
/dts-v1/;
/plugin/;

/ {
    compatible = "brcm,bcm2712";

    fragment@0 {
        target = <&i2s_clk_producer>;
        __overlay__ {
            status = "okay";
        };
    };

    fragment@1 {
        target-path = "/";
        __overlay__ {
            tlv320_dvdd: regulator-tlv320-dvdd {
                compatible = "regulator-fixed";
                regulator-name = "tlv320-dvdd-1v8";
                regulator-min-microvolt = <1800000>;
                regulator-max-microvolt = <1800000>;
                regulator-always-on;
            };

            tlv320_iovdd: regulator-tlv320-iovdd {
                compatible = "regulator-fixed";
                regulator-name = "tlv320-iovdd-3v3";
                regulator-min-microvolt = <3300000>;
                regulator-max-microvolt = <3300000>;
                regulator-always-on;
            };

            tlv320_avdd: regulator-tlv320-avdd {
                compatible = "regulator-fixed";
                regulator-name = "tlv320-avdd-3v3";
                regulator-min-microvolt = <3300000>;
                regulator-max-microvolt = <3300000>;
                regulator-always-on;
            };

            tlv320_hpvdd: regulator-tlv320-hpvdd {
                compatible = "regulator-fixed";
                regulator-name = "tlv320-hpvdd-3v3";
                regulator-min-microvolt = <3300000>;
                regulator-max-microvolt = <3300000>;
                regulator-always-on;
            };

            tlv320_spkvdd: regulator-tlv320-spkvdd {
                compatible = "regulator-fixed";
                regulator-name = "tlv320-spkvdd-5v0";
                regulator-min-microvolt = <5000000>;
                regulator-max-microvolt = <5000000>;
                regulator-always-on;
            };

            /*
             * The T5848 is a control-less I2S microphone from Linux's point
             * of view, so use the generic ICS43432 codec DAI as a simple
             * capture endpoint and keep the speaker path unchanged.
             */
            mic_codec: t5848-codec {
                compatible = "invensense,ics43432";
                #sound-dai-cells = <0>;
            };
        };
    };

    fragment@2 {
        target = <&i2c_arm>;
        __overlay__ {
            #address-cells = <1>;
            #size-cells = <0>;
            status = "okay";

            dac_codec: tlv320dac3100@18 {
                compatible = "ti,tlv320dac3100";
                reg = <0x18>;
                #sound-dai-cells = <0>;
                AVDD-supply = <&tlv320_avdd>;
                IOVDD-supply = <&tlv320_iovdd>;
                DVDD-supply = <&tlv320_dvdd>;
                HPVDD-supply = <&tlv320_hpvdd>;
                SPLVDD-supply = <&tlv320_spkvdd>;
                SPRVDD-supply = <&tlv320_spkvdd>;
                reset-gpios = <&gpio 12 1>; /* GPIO_ACTIVE_LOW */
                ti,pll-clkin-src = <1>;
                status = "okay";
            };
        };
    };

    fragment@3 {
        target-path = "/";
        __overlay__ {
            sound {
                compatible = "simple-audio-card";
                simple-audio-card,name = "barepi-audio";
                simple-audio-card,format = "i2s";
                status = "okay";

                dailink0: simple-audio-card,dai-link@0 {
                    format = "i2s";
                    bitclock-master = <&cpu0>;
                    frame-master = <&cpu0>;

                    cpu0: cpu {
                        sound-dai = <&i2s_clk_producer>;
                    };

                    codec0: codec {
                        sound-dai = <&dac_codec>;
                        system-clock-frequency = <1536000>;
                    };
                };

                dailink1: simple-audio-card,dai-link@1 {
                    format = "i2s";
                    bitclock-master = <&cpu1>;
                    frame-master = <&cpu1>;

                    cpu1: cpu {
                        sound-dai = <&i2s_clk_producer>;
                    };

                    codec1: codec {
                        sound-dai = <&mic_codec>;
                    };
                };
            };
        };
    };
};
DTS_EOF
dtc -@ -I dts -O dtb -o "$OVERLAYS/i2s-audio.dtbo" "$DTS"
rm -f "$DTS"
trap - EXIT
echo "      Installed $OVERLAYS/i2s-audio.dtbo"

# 3. Patch config.txt
echo "[3/7] Configuring config.txt..."
CHANGED=0

if grep -q '^dtparam=i2c_arm=on' "$CONFIG"; then
    echo "      dtparam=i2c_arm=on already set."
elif grep -q '^#.*dtparam=i2c_arm=on' "$CONFIG"; then
    sed -i 's/^#.*dtparam=i2c_arm=on/dtparam=i2c_arm=on/' "$CONFIG"
    CHANGED=1
    echo "      Uncommented dtparam=i2c_arm=on."
else
    if grep -q '^dtparam=audio=on' "$CONFIG"; then
        sed -i '/^dtparam=audio=on/a dtparam=i2c_arm=on' "$CONFIG"
    else
        echo "dtparam=i2c_arm=on" >> "$CONFIG"
    fi
    CHANGED=1
    echo "      Added dtparam=i2c_arm=on."
fi

if grep -q '^dtparam=i2s=on' "$CONFIG"; then
    echo "      dtparam=i2s=on already set."
elif grep -q '^#.*dtparam=i2s=on' "$CONFIG"; then
    sed -i 's/^#.*dtparam=i2s=on/dtparam=i2s=on/' "$CONFIG"
    CHANGED=1
    echo "      Uncommented dtparam=i2s=on."
else
    if grep -q '^dtparam=audio=on' "$CONFIG"; then
        sed -i '/^dtparam=audio=on/a dtparam=i2s=on' "$CONFIG"
    else
        echo "dtparam=i2s=on" >> "$CONFIG"
    fi
    CHANGED=1
    echo "      Added dtparam=i2s=on."
fi

if grep -q '^dtoverlay=i2s-audio' "$CONFIG"; then
    echo "      dtoverlay=i2s-audio already set."
elif grep -q '^#.*dtoverlay=i2s-audio' "$CONFIG"; then
    sed -i 's/^#.*dtoverlay=i2s-audio/dtoverlay=i2s-audio/' "$CONFIG"
    CHANGED=1
    echo "      Uncommented dtoverlay=i2s-audio."
else
    if grep -q '^dtparam=i2s=on' "$CONFIG"; then
        sed -i '/^dtparam=i2s=on/a dtoverlay=i2s-audio' "$CONFIG"
    else
        echo "dtoverlay=i2s-audio" >> "$CONFIG"
    fi
    CHANGED=1
    echo "      Added dtoverlay=i2s-audio."
fi

if grep -q '^gpio=12=op,dh' "$CONFIG"; then
    echo "      gpio=12=op,dh already set."
elif grep -q '^#.*gpio=12=op,dh' "$CONFIG"; then
    sed -i 's/^#.*gpio=12=op,dh/gpio=12=op,dh/' "$CONFIG"
    CHANGED=1
    echo "      Uncommented gpio=12=op,dh."
elif grep -q '^gpio=12=' "$CONFIG"; then
    echo "      Existing gpio=12=... line found; left unchanged." >&2
    echo "      Make sure it keeps the DAC out of reset (gpio=12=op,dh)." >&2
else
    if grep -q '^dtoverlay=i2s-audio' "$CONFIG"; then
        sed -i '/^dtoverlay=i2s-audio/a gpio=12=op,dh' "$CONFIG"
    elif grep -q '^dtparam=i2s=on' "$CONFIG"; then
        sed -i '/^dtparam=i2s=on/a gpio=12=op,dh' "$CONFIG"
    else
        echo "gpio=12=op,dh" >> "$CONFIG"
    fi
    CHANGED=1
    echo "      Added gpio=12=op,dh."
fi

# 4. Install .asoundrc
echo "[4/7] Installing .asoundrc..."
cat > "$TARGET_HOME/.asoundrc" << 'ASOUNDRC'
defaults.pcm.rate_converter "samplerate_best"

pcm.!default {
    type plug
    slave {
        pcm "filtered"
        rate 48000
    }
}

pcm.filtered {
    type ladspa
    slave.pcm "out"
    path "/usr/lib/ladspa"
    plugins [
        {
            label highpass_iir
            input {
                controls [ 100 1 ]
            }
        }
        {
            label hardLimiter
            input {
                controls [ 0 1 0 ]
            }
        }
    ]
}

pcm.out {
    type plug
    slave {
        pcm "hw:CARD=barepiaudio,DEV=1"
        format S16_LE
    }
}

pcm.barepi_playback {
    type plug
    slave {
        pcm "out"
        rate 48000
    }
}

pcm.barepi_capture {
    type plug
    slave {
        pcm "hw:CARD=barepiaudio,DEV=0"
        format S32_LE
        rate 48000
        channels 2
    }
}

ctl.!default {
    type hw
    card barepiaudio
}
ASOUNDRC
chown "$TARGET_USER:$TARGET_USER" "$TARGET_HOME/.asoundrc"
echo "      Installed $TARGET_HOME/.asoundrc"

# 5. Install volume.sh
echo "[5/7] Installing volume.sh..."
cat > "$TARGET_HOME/volume.sh" << 'VOLSCRIPT'
#!/bin/bash
# DAC3100 speaker-path helper
# Usage: ./volume.sh [dac] [driver] [analog]
#   dac:     0-175  (digital volume, 127=0dB)
#   driver:  0-3    (class-D gain: 6/12/18/24 dB)
#   analog:  0-127  (analog volume, 127=0dB)

set -euo pipefail

CARD_ID="barepiaudio"
CARD_NAME="barepi-audio"
DEFAULT_DAC=127
DEFAULT_DRIVER=0
DEFAULT_ANALOG=50

find_card() {
  awk -v id="$CARD_ID" -v name="$CARD_NAME" '
    index($0, "[" id "]") || index($0, name) { print $1; exit }
  ' /proc/asound/cards
}

show_status() {
  echo "Using ALSA card: $CARD ($CARD_NAME)"
  amixer -c "$CARD" sget 'DAC' || true
  amixer -c "$CARD" sget 'DAC Left Input' || true
  amixer -c "$CARD" sget 'DAC Right Input' || true
  amixer -c "$CARD" sget 'Output Left From Left DAC' || true
  amixer -c "$CARD" sget 'Output Right From Right DAC' || true
  amixer -c "$CARD" sget 'Speaker' || true
  amixer -c "$CARD" sget 'Speaker Driver' || true
  amixer -c "$CARD" sget 'Speaker Analog' || true
}

enable_speaker_path() {
  echo "Enabling TLV320 speaker path..."
  amixer -c "$CARD" sset 'DAC Left Input' 'Left Data' >/dev/null || true
  amixer -c "$CARD" sset 'DAC Right Input' 'Right Data' >/dev/null || true
  amixer -c "$CARD" sset 'Output Left From Left DAC' on >/dev/null || true
  amixer -c "$CARD" sset 'Output Right From Right DAC' on >/dev/null || true
  amixer -c "$CARD" sset 'Speaker' on >/dev/null || true
  amixer -c "$CARD" sset 'HP Left' off >/dev/null || true
  amixer -c "$CARD" sset 'HP Right' off >/dev/null || true
  amixer -c "$CARD" sset 'HP Driver' 0 >/dev/null || true
  amixer -c "$CARD" sset 'HP Driver' off >/dev/null || true
}

CARD="$(find_card || true)"
if [ -z "$CARD" ]; then
  echo "Could not find ALSA card for $CARD_NAME." >&2
  echo "Reboot after setup, then check /proc/asound/cards." >&2
  exit 1
fi

if [ $# -eq 0 ]; then
  show_status
  exit 0
fi

DAC=${1:-$DEFAULT_DAC}
DRIVER=${2:-$DEFAULT_DRIVER}
ANALOG=${3:-$DEFAULT_ANALOG}

enable_speaker_path

echo "Setting levels..."
amixer -c "$CARD" sset 'DAC' "${DAC},${DAC}" >/dev/null
amixer -c "$CARD" sset 'Speaker Driver' "$DRIVER" >/dev/null || true
amixer -c "$CARD" sset 'Speaker Driver' on >/dev/null || true
amixer -c "$CARD" sset 'Speaker Analog' "$ANALOG" >/dev/null || true

alsactl store "$CARD" >/dev/null 2>&1 || alsactl store >/dev/null 2>&1 || true

echo ""
echo "Using ALSA card: $CARD ($CARD_NAME)"
echo "Set: dac=$DAC driver=$DRIVER analog=$ANALOG"
echo "Example: ./volume.sh 127 0 50"
VOLSCRIPT
chown "$TARGET_USER:$TARGET_USER" "$TARGET_HOME/volume.sh"
chmod +x "$TARGET_HOME/volume.sh"
echo "      Installed $TARGET_HOME/volume.sh"

# 6. Install a mic test helper
echo "[6/7] Installing test-mics.sh..."
cat > "$TARGET_HOME/test-mics.sh" << 'MICTEST'
#!/bin/bash
# Capture a short stereo sample from the barepi mic path and print per-channel stats.
# Usage:
#   ./test-mics.sh [seconds] [out.wav]
#   ./test-mics.sh --thsel-high [seconds] [out.wav]
#   ./test-mics.sh --thsel-low  [seconds] [out.wav]

set -euo pipefail

RATE=48000
FORMAT=S32_LE
CHANNELS=2
DURATION=3
OUT=""
THSEL=""
CAPTURE_PCM="${CAPTURE_PCM:-barepi_capture}"

usage() {
  cat <<'EOF'
Usage:
  ./test-mics.sh [seconds] [out.wav]
  ./test-mics.sh --thsel-high [seconds] [out.wav]
  ./test-mics.sh --thsel-low  [seconds] [out.wav]
EOF
}

set_thsel() {
  local level="$1"
  command -v pinctrl >/dev/null 2>&1 || return 0
  if [ "$level" = "high" ]; then
    sudo pinctrl set 8 op pn dh >/dev/null
  else
    sudo pinctrl set 8 op pn dl >/dev/null
  fi
}

prepare_wake_pin() {
  command -v pinctrl >/dev/null 2>&1 || return 0
  sudo pinctrl set 13 ip pn >/dev/null 2>&1 || true
}

show_pins() {
  command -v pinctrl >/dev/null 2>&1 || return 0
  sudo pinctrl get 8,13 || true
}

positional=()
while [ $# -gt 0 ]; do
  case "$1" in
    --thsel-high)
      THSEL="high"
      shift
      ;;
    --thsel-low)
      THSEL="low"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      positional+=("$1")
      shift
      ;;
  esac
done

if [ ${#positional[@]} -ge 1 ]; then
  DURATION="${positional[0]}"
fi
if [ ${#positional[@]} -ge 2 ]; then
  OUT="${positional[1]}"
fi
if [ -z "$OUT" ]; then
  OUT="/tmp/barepi-mics-$(date +%Y%m%d-%H%M%S).wav"
fi

if ! arecord -L 2>/dev/null | grep -qx 'barepi_capture'; then
  CAPTURE_PCM="plughw:CARD=barepiaudio,DEV=0"
fi

if [ -n "$THSEL" ]; then
  set_thsel "$THSEL"
fi
prepare_wake_pin
show_pins

echo "Recording ${DURATION}s from ${CAPTURE_PCM} -> ${OUT}"
arecord -q -D "$CAPTURE_PCM" -f "$FORMAT" -c "$CHANNELS" -r "$RATE" -d "$DURATION" "$OUT"

python3 - <<'PY' "$OUT"
import math
import struct
import sys
import wave

path = sys.argv[1]
with wave.open(path, 'rb') as w:
    channels = w.getnchannels()
    width = w.getsampwidth()
    rate = w.getframerate()
    frames = w.getnframes()
    data = w.readframes(frames)

if channels != 2 or width != 4:
    print(f"Recorded {channels}ch / {width * 8}bit at {rate} Hz")
    sys.exit(0)

samples = list(struct.iter_unpack('<ii', data))
left = [l for l, _ in samples]
right = [r for _, r in samples]

def rms(xs):
    return math.sqrt(sum(x * x for x in xs) / len(xs)) if xs else 0.0

def peak(xs):
    return max((abs(x) for x in xs), default=0)

l_rms = rms(left)
r_rms = rms(right)
l_peak = peak(left)
r_peak = peak(right)
full_scale = float(2**31 - 1)

print(f"Saved: {path}")
print(f"Rate: {rate} Hz  Frames: {frames}")
print(f"Left : rms={int(l_rms):>10} peak={l_peak:>10} peak_dBFS={20*math.log10(max(l_peak,1)/full_scale):6.1f}")
print(f"Right: rms={int(r_rms):>10} peak={r_peak:>10} peak_dBFS={20*math.log10(max(r_peak,1)/full_scale):6.1f}")
PY
MICTEST
chown "$TARGET_USER:$TARGET_USER" "$TARGET_HOME/test-mics.sh"
chmod +x "$TARGET_HOME/test-mics.sh"
echo "      Installed $TARGET_HOME/test-mics.sh"

# 7. Install a boot-time restore helper so the speaker path comes up after reboot
echo "[7/7] Installing boot-time mixer restore..."
cat > "$RESTORE_HELPER" << 'RESTORESCRIPT'
#!/bin/bash
set -euo pipefail

CARD_ID="barepiaudio"
CARD_NAME="barepi-audio"
I2C_BUS="${I2C_BUS:-1}"
CODEC_ADDR="${CODEC_ADDR:-0x18}"
MIC_THSEL_LEVEL="${MIC_THSEL_LEVEL:-high}"
DAC=${1:-127}
DRIVER=${2:-0}
ANALOG=${3:-50}

find_card() {
  awk -v id="$CARD_ID" -v name="$CARD_NAME" '
    index($0, "[" id "]") || index($0, name) { print $1; exit }
  ' /proc/asound/cards
}

codec_ready() {
  i2cget -y -f "$I2C_BUS" "$CODEC_ADDR" 0 >/dev/null 2>&1
}

set_mixer() {
  amixer -q -c "$CARD" "$@" >/dev/null 2>&1 || true
}

set_mic_pins() {
  if ! command -v pinctrl >/dev/null 2>&1; then
    return 0
  fi

  case "$MIC_THSEL_LEVEL" in
    low)
      pinctrl set 8 op pn dl >/dev/null 2>&1 || true
      ;;
    *)
      pinctrl set 8 op pn dh >/dev/null 2>&1 || true
      ;;
  esac

  pinctrl set 13 ip pn >/dev/null 2>&1 || true
}

CARD=""
for _ in $(seq 1 20); do
  CARD="$(find_card || true)"
  if [ -n "$CARD" ]; then
    break
  fi
  sleep 1
done

if [ -z "$CARD" ]; then
  echo "barepi-audio card not present; skipping restore."
  exit 0
fi

for _ in $(seq 1 20); do
  if codec_ready; then
    break
  fi
  sleep 1
done

if ! codec_ready; then
  echo "TLV320 codec not responding on i2c-${I2C_BUS} addr ${CODEC_ADDR}; skipping restore."
  exit 0
fi

set_mic_pins
set_mixer sset 'DAC Left Input' 'Left Data'
set_mixer sset 'DAC Right Input' 'Right Data'
set_mixer sset 'Output Left From Left DAC' on
set_mixer sset 'Output Right From Right DAC' on
set_mixer sset 'Speaker' on
set_mixer sset 'Speaker Driver' "$DRIVER"
set_mixer sset 'Speaker Driver' on
set_mixer sset 'Speaker Analog' "$ANALOG"
set_mixer sset 'DAC' "${DAC},${DAC}"
set_mixer sset 'HP Left' off
set_mixer sset 'HP Right' off
set_mixer sset 'HP Driver' '0,0'
set_mixer sset 'HP Driver' off
RESTORESCRIPT
chmod 755 "$RESTORE_HELPER"

cat > "$RESTORE_SERVICE" << EOF
[Unit]
Description=Restore barepi TLV320 speaker routing
After=sound.target
Wants=sound.target

[Service]
Type=oneshot
ExecStart=$RESTORE_HELPER 127 0 50

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable barepi-audio-restore.service >/dev/null

if grep -q '\[barepiaudio\]' /proc/asound/cards 2>/dev/null || grep -q 'barepi-audio' /proc/asound/cards 2>/dev/null; then
    "$RESTORE_HELPER" 127 0 50 || true
    echo "      Applied mixer defaults to the live barepi-audio card."
else
    echo "      Live barepi-audio card not present yet; defaults will be applied at boot."
fi

echo ""
echo "=== Done ==="
if [ "$CHANGED" -eq 1 ]; then
    echo "config.txt was modified — reboot to load the overlay/reset config."
else
    echo "config.txt unchanged. Reboot if this is a first-time setup."
fi
echo "Speaker routing defaults and mic GPIO defaults will be restored automatically after boot."
echo "Manual speaker check/adjust: $TARGET_HOME/volume.sh"
echo "Manual mic test: $TARGET_HOME/test-mics.sh"