#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ENV_FILE="${ENV_FILE:-$ROOT_DIR/.env.pp-chatocr.verify}"
SAMPLE_FILE="${SAMPLE_FILE:-$ROOT_DIR/samples/images/56CDF923B8CAE7000FE3274C10ECBC0C.jpg}"

load_env() {
  while IFS= read -r line || [[ -n "$line" ]]; do
    [[ -z "$line" || "$line" =~ ^[[:space:]]*# ]] && continue
    [[ "$line" != *=* ]] && continue
    local key="${line%%=*}"
    local value="${line#*=}"
    key="$(echo "$key" | xargs)"
    [[ -z "$key" ]] && continue
    export "$key=$value"
  done < "$ENV_FILE"
}

if [[ ! -f "$ENV_FILE" ]]; then
  echo "Missing env file: $ENV_FILE" >&2
  echo "Copy .env.pp-chatocr.verify.example to .env.pp-chatocr.verify and fill local secrets." >&2
  exit 1
fi

load_env

echo "== PP-ChatOCR + LLM verification =="
echo "env file: $ENV_FILE"
echo "sample: $SAMPLE_FILE"
echo "visual endpoint: $PP_CHAT_OCR_VISUAL_ENDPOINT"
echo "chat endpoint: $PP_CHAT_OCR_CHAT_ENDPOINT"
echo "llm model id: ${LLM_MODEL_ID:-MiniMax-M2.7}"

python3 - "$SAMPLE_FILE" <<'PY'
import base64
import json
import os
import sys
import urllib.error
import urllib.request

sample_file = sys.argv[1]
visual_endpoint = os.environ["PP_CHAT_OCR_VISUAL_ENDPOINT"]
chat_endpoint = os.environ["PP_CHAT_OCR_CHAT_ENDPOINT"]
settings_path = os.environ.get("MAODOU_MODEL_SETTINGS", "/Users/alphaPlanet/.maodou/settings/model-settings.json")
model_id = os.environ.get("LLM_MODEL_ID", "MiniMax-M2.7")
chat_hit_required = os.environ.get("CHAT_HIT_REQUIRED", "false").lower() == "true"

raw_keys = os.environ.get("PP_CHAT_OCR_EXTRACTION_KEYS", "")
key_list = [item.strip() for item in raw_keys.split(",") if item.strip()]
if not key_list:
    key_list = ["经营者名称", "申请日期", "经营场所地址", "证件号码", "联系电话"]

def post_json(endpoint, payload, timeout=120):
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    request = urllib.request.Request(
        endpoint,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise SystemExit(f"HTTP {exc.code} from {endpoint}: {detail[:1000]}")

with open(sample_file, "rb") as fh:
    encoded = base64.b64encode(fh.read()).decode("ascii")

visual_payload = {
    "file": encoded,
    "fileType": 1,
    "useDocOrientationClassify": True,
    "useDocUnwarping": True,
    "useTextlineOrientation": True,
    "useSealRecognition": True,
    "useTableRecognition": True,
}

visual = post_json(visual_endpoint, visual_payload)
visual_ok = visual.get("errorCode") == 0
visual_info = visual.get("result", {}).get("visualInfo", [])
layout_results = visual.get("result", {}).get("layoutParsingResults", [])
print(json.dumps({
    "stage": "visual",
    "serviceReachable": True,
    "success": visual_ok,
    "errorCode": visual.get("errorCode"),
    "errorMsg": visual.get("errorMsg", ""),
    "visualInfoCount": len(visual_info),
    "layoutPages": len(layout_results),
}, ensure_ascii=False))
if not visual_ok:
    raise SystemExit("PP-ChatOCR visual service returned non-zero errorCode")

settings = json.load(open(settings_path))
model = settings.get("models", {}).get(model_id)
if not model:
    raise SystemExit(f"LLM model {model_id} not found in {settings_path}")

chat_payload = {
    "keyList": key_list,
    "visualInfo": visual_info,
    "useVectorRetrieval": False,
    "minCharacters": 500,
    "textTaskDescription": "请从文档 OCR 结果中抽取 keyList 指定字段的值。字段没有出现时返回空字符串，不要编造。",
    "textOutputFormat": "返回 JSON 对象，键名必须来自 keyList，值为文档中对应的原文。",
    "textRulesStr": "保留日期、编号、地址、姓名、证件号码等原文格式；复选框或单选框字段只返回被勾选或最明确的选项；无法确认时返回空字符串。",
    "tableTaskDescription": "请优先从表格单元格的字段和值关系中抽取 keyList 对应内容。",
    "tableOutputFormat": "返回 JSON 对象，键名必须来自 keyList，值为表格中对应的原文。",
    "tableRulesStr": "相邻单元格、同一行或同一区块中与字段标签对应的内容视为候选值；不要把字段标签本身当作值。",
    "chatBotConfig": {
        "module_name": "chat_bot",
        "model_name": model_id,
        "base_url": model.get("baseUrl", ""),
        "api_type": os.environ.get("LLM_API_TYPE", "openai"),
        "api_key": model.get("apiKey", ""),
        "enable_thinking": os.environ.get("LLM_THINKING_ENABLED", "false").lower() == "true",
    },
}

chat = post_json(chat_endpoint, chat_payload)
chat_ok = chat.get("errorCode") == 0
chat_result = chat.get("result", {}).get("chatResult")

hit_count = 0
if isinstance(chat_result, dict):
    for value in chat_result.values():
        if value is not None and str(value).strip():
            hit_count += 1
elif isinstance(chat_result, list):
    hit_count = sum(1 for item in chat_result if item)

print(json.dumps({
    "stage": "chat",
    "serviceReachable": True,
    "success": chat_ok,
    "errorCode": chat.get("errorCode"),
    "errorMsg": chat.get("errorMsg", ""),
    "requestedKeyCount": len(key_list),
    "fieldHitCount": hit_count,
    "fieldHit": hit_count > 0,
}, ensure_ascii=False))

if not chat_ok:
    raise SystemExit("PP-ChatOCR chat service returned non-zero errorCode")
if chat_hit_required and hit_count <= 0:
    raise SystemExit("PP-ChatOCR chat service is reachable, but no requested fields were hit")
PY

echo
echo "PP-ChatOCR verification completed."