#!/bin/bash

set -x

export HOME=/root
export USER=root
export DEBIAN_FRONTEND=noninteractive
export NEEDRESTART_MODE=a
export UCF_FORCE_CONFOLD=1	
export LOCK_FILE="/var/lib/dpkg/lock-frontend"
export LOCK_FILE2="/var/lib/apt/lists/lock"
export MAX_WAIT=300
export WAIT_INTERVAL=10
export ELAPSED=0

if systemctl is-active --quiet unattended-upgrades; then
    systemctl stop unattended-upgrades
fi

while fuser $LOCK_FILE $LOCK_FILE2 >/dev/null 2>&1; do
    sleep $WAIT_INTERVAL
    ELAPSED=$((ELAPSED + WAIT_INTERVAL))
    if [ $ELAPSED -ge $MAX_WAIT ]; then
        PIDS=$(fuser $LOCK_FILE $LOCK_FILE2 2>/dev/null)
        if [ -n "$PIDS" ]; then
            kill -9 $PIDS
        fi
        dpkg --configure -a
        break
    fi
done
apt update
apt install -y git wget curl vim unzip net-tools build-essential cmake fonts-wqy-zenhei
apt install -y python3-venv python3-pip

#wget https://download.cloudclusters.io/cuda_11.7_linux.run
wget -q https://download-nocix.cloudclusters.io/cuda_11.7_linux.run
#wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run
#sh cuda_11.7.0_515.43.04_linux.run --silent --toolkit --override
sh cuda_11.7_linux.run --silent --toolkit --override
#wget https://download.cloudclusters.io/cudnn_ubuntu22_amd64.deb
wget -q https://download-nocix.cloudclusters.io/cudnn_ubuntu22_amd64.deb
dpkg -i cudnn_ubuntu22_amd64.deb
cp /var/cudnn-local-repo-ubuntu2204-8.9.7.29/cudnn-local-*-keyring.gpg /usr/share/keyrings/
apt update
apt install libcudnn8 libcudnn8-dev
echo 'export PATH=/usr/local/cuda-11.7/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
rm -fr cuda_11.7_linux.run cudnn_ubuntu22_amd64.deb

pip install paddlepaddle-gpu==2.6.2 -f https://www.paddlepaddle.org.cn/whl/windows/mkl/avx/stable.html
python3 -c "import paddle; print(paddle.device.is_compiled_with_cuda()); print(paddle.get_device())"
pip install paddleocr==2.7.0.3
pip install gradio
pip uninstall -y numpy opencv-python opencv-python-headless
pip install numpy==1.26.4
pip install opencv-python==4.6.0.66
pip install opencv-python-headless==4.9.0.80

cat << 'EOF' > /root/.web.py
import gradio as gr
from paddleocr import PaddleOCR, draw_ocr
from PIL import Image, ImageFont
import numpy as np
import os
import glob

def get_chinese_font():
    font_paths = [
        "/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc",
        "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",
        "/usr/share/fonts/truetype/arphic/ukai.ttc",
        "/usr/share/fonts/truetype/arphic/uming.ttc",
        "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",
        "simsun.ttc", 
    ]
    for path in font_paths:
        if os.path.exists(path):
            return path
    return None

font_path = get_chinese_font()
if font_path:
    print(f"[INFO] Using font: {font_path}")
else:
    print("[WARNING] No Chinese font found, boxes only, text will not render.")

ocr = PaddleOCR(use_angle_cls=True, lang='ch') 
def ocr_fn(img):
    if not isinstance(img, np.ndarray):
        img = np.array(img)

    result = ocr.ocr(img)

    if not result or not result[0]:
        return img, "No text detected"

    boxes = [line[0] for line in result[0]]
    texts = [line[1][0] for line in result[0]]
    scores = [line[1][1] for line in result[0]]

    if font_path:
        img_result = draw_ocr(img, boxes, texts, scores, font_path=font_path)
    else:
        img_result = draw_ocr(img, boxes, texts, scores, font_path=None)

    img_result = Image.fromarray(img_result)
    text_output = "\n".join([f"{t} ({s:.2f})" for t, s in zip(texts, scores)])
    return img_result, text_output

demo = gr.Interface(
    fn=ocr_fn,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=[gr.Image(label="OCR Result"), gr.Textbox(label="Detected Text")],
    title="PaddleOCR Gradio Demo",
    description="Upload an image to perform OCR. Supports Chinese and English."
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=5000)
EOF

apt install supervisor

cat << 'EOF' > /etc/supervisor/conf.d/web.conf
[program:web]
command=/usr/bin/python3 /root/.web.py
environment=PATH="/usr/local/cuda-11.7/bin:$PATH",LD_LIBRARY_PATH="/usr/local/cuda-11.7/lib64:/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH"
autostart=true
autorestart=true
stopasgroup=true
killasgroup=true
stderr_logfile=/var/log/supervisor/web_error.log
stdout_logfile=/var/log/supervisor/web_stdout.log
EOF

supervisorctl update

OPEN_PORT=5000
TIMEOUT=600
INTERVAL=10
ELAPSED=0
echo "Checking if port $OPEN_PORT is listening (timeout: $TIMEOUT seconds)..."
while [ $ELAPSED -lt $TIMEOUT ]; do
	if ss -tulnp | grep -q ":$OPEN_PORT"; then
		echo "Port $OPEN_PORT is listening (after ${ELAPSED}s)"
		exit 0
	fi
	sleep $INTERVAL
	ELAPSED=$((ELAPSED + INTERVAL))
	echo "Waited ${ELAPSED}s..."
done
echo "Port $OPEN_PORT did not start listening within $TIMEOUT seconds, exiting"
exit 1
