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

CONTAINER_NAME="${RAGFLOW_CONTAINER_NAME:-ffoa-dev-ragflow}"
ADMIN_EMAIL="${RAGFLOW_ADMIN_EMAIL:-}"
ADMIN_PASSWORD="${RAGFLOW_ADMIN_PASSWORD:-}"
ADMIN_NICKNAME="${RAGFLOW_ADMIN_NICKNAME:-RAGFlow Admin}"
API_TOKEN="${RAGFLOW_API_TOKEN:-}"
DASHSCOPE_KEY="${RAGFLOW_DASHSCOPE_API_KEY:-${DASHSCOPE_API_KEY:-}}"

if [[ -z "${ADMIN_EMAIL}" || -z "${ADMIN_PASSWORD}" ]]; then
  echo "缺少必填环境变量：RAGFLOW_ADMIN_EMAIL、RAGFLOW_ADMIN_PASSWORD"
  exit 1
fi

docker exec -i \
  -e RAGFLOW_ADMIN_EMAIL="${ADMIN_EMAIL}" \
  -e RAGFLOW_ADMIN_PASSWORD="${ADMIN_PASSWORD}" \
  -e RAGFLOW_ADMIN_NICKNAME="${ADMIN_NICKNAME}" \
  -e RAGFLOW_API_TOKEN="${API_TOKEN}" \
  -e RAGFLOW_DASHSCOPE_API_KEY="${DASHSCOPE_KEY}" \
  "${CONTAINER_NAME}" python - <<'PY'
import os
import sys

admin_email = os.environ.get("RAGFLOW_ADMIN_EMAIL")
admin_password = os.environ.get("RAGFLOW_ADMIN_PASSWORD")
admin_nickname = os.environ.get("RAGFLOW_ADMIN_NICKNAME") or "RAGFlow Admin"
api_token = os.environ.get("RAGFLOW_API_TOKEN")
dashscope_key = os.environ.get("RAGFLOW_DASHSCOPE_API_KEY")

if not admin_email or not admin_password:
  print("缺少必填环境变量：RAGFLOW_ADMIN_EMAIL、RAGFLOW_ADMIN_PASSWORD", file=sys.stderr)
  sys.exit(1)

from api.db import UserTenantRole
from api.db.services.user_service import UserService, TenantService, UserTenantService
from api.db.services.tenant_llm_service import TenantLLMService
from api.db.services.api_service import APITokenService
from api.db.services.llm_service import get_init_tenant_llm
from api.utils.api_utils import generate_confirmation_token
from common.constants import StatusEnum
from common.misc_utils import get_uuid
from common import settings

# Ensure settings are initialized from service_conf/env
settings.init_settings()

existing = UserService.query_user_by_email(admin_email) or []
user = existing[0] if len(existing) else None
if user is None:
  user_id = get_uuid()
  user = UserService.save(
    id=user_id,
    email=admin_email,
    nickname=admin_nickname,
    password=admin_password,
    is_superuser=False,
    status=StatusEnum.VALID.value,
    login_channel="password",
  )
  print("已创建 RAGFlow 账号：%s" % admin_email)
else:
  user_id = user.id
  print("RAGFlow 账号已存在：%s" % admin_email)

parser_ids = settings.PARSERS or (
  "naive:General,qa:Q&A,resume:Resume,manual:Manual,table:Table,paper:Paper,"
  "book:Book,laws:Laws,presentation:Presentation,picture:Picture,one:One,"
  "audio:Audio,email:Email,tag:Tag"
)

tenants = TenantService.query(id=user_id) or []
if len(tenants) == 0:
  TenantService.insert(
    id=user_id,
    name=f"{admin_nickname}‘s Kingdom",
    llm_id=settings.CHAT_MDL,
    embd_id=settings.EMBEDDING_MDL,
    asr_id=settings.ASR_MDL,
    img2txt_id=settings.IMAGE2TEXT_MDL,
    parser_ids=parser_ids,
  )

user_tenants = UserTenantService.query(user_id=user_id, tenant_id=user_id) or []
if len(user_tenants) == 0:
  UserTenantService.insert(
    tenant_id=user_id,
    user_id=user_id,
    role=UserTenantRole.OWNER.value,
    invited_by=user_id,
    status=StatusEnum.VALID.value,
  )

try:
  tenant_llm = get_init_tenant_llm(user_id)
  if tenant_llm:
    TenantLLMService.insert_many(tenant_llm)
except Exception as exc:
  print("初始化 Tenant LLM 失败，已跳过：%s" % exc)

if dashscope_key:
  try:
    updated = (
      TenantLLMService.model
      .update(
        api_key=dashscope_key,
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
      )
      .where(
        (TenantLLMService.model.tenant_id == user_id) &
        (TenantLLMService.model.llm_factory == "Tongyi-Qianwen") &
        (TenantLLMService.model.model_type == "embedding") &
        ((TenantLLMService.model.api_key.is_null()) |
         (TenantLLMService.model.api_key == "") |
         (TenantLLMService.model.api_key == "dev-dashscope-api-key"))
      )
      .execute()
    )
    if updated:
      print(f"已更新通义 embedding api_key: {updated} 条")
  except Exception as exc:
    print(f"更新通义 embedding api_key 失败，已跳过：{exc}")

raw_token = api_token or generate_confirmation_token()
if raw_token.startswith("ragflow-"):
  raw_token = raw_token[len("ragflow-"):]

try:
  if api_token and api_token.startswith("ragflow-"):
    existing_raw = APITokenService.query(tenant_id=user_id, token=raw_token) or []
    existing_prefixed = APITokenService.query(tenant_id=user_id, token=api_token) or []
    if existing_prefixed and not existing_raw:
      APITokenService.model.update(token=raw_token).where(
        (APITokenService.model.tenant_id == user_id) &
        (APITokenService.model.token == api_token)
      ).execute()
    elif existing_prefixed and existing_raw:
      APITokenService.model.delete().where(
        (APITokenService.model.tenant_id == user_id) &
        (APITokenService.model.token == api_token)
      ).execute()
except Exception as exc:
  print(f"规范化 RAGFlow token 失败，已跳过：{exc}")

existing_tokens = APITokenService.query(tenant_id=user_id, token=raw_token) or []
if len(existing_tokens) == 0:
  APITokenService.save(
    tenant_id=user_id,
    token=raw_token,
    dialog_id=None,
    source="none",
    beta=None,
  )
print("RAGFLOW_API_KEY=%s" % raw_token)
PY