src/storage.py

# storage.py  (Hugging Face Dataset 버전)
# ───────────────────────────────────────────────────────────────────────────
# 데이터 저장 — SQLite도 구글시트도 쓰지 않고 '이중 안전망'으로 저장한다.
#
#   안전망 1) HF Dataset 업로드 : 영구 저장 + 여러 참여자 통합 (구글시트 대체)
#   안전망 2) 메모리(session_state): 사이드바에서 즉시 CSV 다운로드 (백업)
#
# 왜 HF Dataset인가:
#   - 구글 방식(서비스계정·JSON 키·시트 공유)이 초보자에게 너무 복잡하다.
#   - 이미 HF Spaces에 배포하므로, 같은 HF 안에서 '토큰 1개'만 더 만들면 끝.
#   - Dataset 저장소는 Space와 달리 재시작에 영향받지 않는 git 저장소라
#     데이터가 영구 보존된다(Space가 슬립/재빌드돼도 안전).
#
# 작동 방식:
#   - 모든 기록은 일단 메모리 리스트(rows_*)에 쌓인다(안전망 2).
#   - '참여자 완료' 시점에만, 그때까지 모인 전체 CSV를 Dataset에 덮어쓴다(안전망 1).
#     → 매 메시지마다 올리면 너무 잦으므로, 완료 시 한 번만 통째로 올린다.
#   - 단, 메모리는 '현재 브라우저 세션'만 담는다. 여러 참여자를 한 파일로
#     합치기 위해, 업로드 전에 Dataset의 기존 CSV를 내려받아 '머지(merge)'한다.
#
# 설정(1회): README/강의 7장 참조.
#   1) HF에서 빈 Dataset 생성 (예: yourname/career-chatbot-data, Private 권장)
#   2) write 권한 토큰 발급 → Space secret에 HF_TOKEN으로 등록
#   3) 아래 DATASET_REPO를 본인 것으로 교체
# HF_TOKEN이 없으면 업로드는 조용히 건너뛰고 사이드바 다운로드만 동작한다
# (로컬 개발 시 편리).
# ───────────────────────────────────────────────────────────────────────────
import os
import io
from datetime import datetime, timezone

import pandas as pd

# huggingface_hub은 Dataset 업로드를 쓸 때만 필요. 없으면 메모리 모드로만 동작.
try:
    from huggingface_hub import HfApi, hf_hub_download
    _HAS_HF = True
except ImportError:
    _HAS_HF = False

# ▼▼▼ 본인이 만든 Dataset 저장소 이름으로 교체 ▼▼▼
DATASET_REPO = "jonghhhh/career-chatbot-data"
# ▲▲▲ (HF에서 New Dataset으로 만든 'yourname/저장소이름') ▲▲▲

HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN_WRITE")  # Space secret 또는 .env에 넣은 write 토큰

# Dataset 안에 저장될 3개 CSV 파일 이름
FILES = {
    "participants": "participants.csv",
    "messages": "messages.csv",
    "surveys": "surveys.csv",
}


def now_iso():
    """UTC ISO 8601 타임스탬프 (예: 2026-05-27T12:34:56+00:00)."""
    return datetime.now(timezone.utc).isoformat()


# ── 메모리(안전망 2): 모든 기록을 세션 리스트에 누적 ─────────────────────────
def init_state(state):
    """세션 시작 시 누적 리스트 3개를 준비(이미 있으면 그대로)."""
    state.setdefault("rows_participants", [])
    state.setdefault("rows_messages", [])
    state.setdefault("rows_surveys", [])


def add_participant(state, pid, condition):
    """참여자 배정 기록(메모리에만 추가. 업로드는 완료 시 일괄)."""
    state["rows_participants"].append(
        {"participant_id": pid, "condition": condition,
         "assigned_at": now_iso(), "completed": 0})


def add_message(state, pid, condition, turn, role, content, latency_ms=""):
    """대화 한 턴 기록(메모리)."""
    state["rows_messages"].append(
        {"participant_id": pid, "condition": condition, "turn": turn,
         "role": role, "content": content, "latency_ms": latency_ms,
         "created_at": now_iso()})


def add_survey(state, pid, condition, phase, question_id, answer):
    """설문 한 문항 기록(메모리)."""
    state["rows_surveys"].append(
        {"participant_id": pid, "condition": condition, "phase": phase,
         "question_id": question_id, "answer": str(answer),
         "created_at": now_iso()})


# ── HF Dataset(안전망 1): 기존 데이터와 머지 후 업로드 ───────────────────────
def _download_existing(api, filename):
    """Dataset에 이미 있는 CSV를 DataFrame으로 읽어온다. 없으면 빈 DF."""
    try:
        path = hf_hub_download(
            repo_id=DATASET_REPO, filename=filename,
            repo_type="dataset", token=HF_TOKEN,
        )
        return pd.read_csv(path)
    except Exception:
        # 파일이 아직 없거나(첫 업로드) 네트워크 문제 → 빈 DF로 시작
        return pd.DataFrame()


def _upload_csv(api, df, filename):
    """DataFrame을 CSV 바이트로 만들어 Dataset에 덮어쓰기 업로드."""
    buf = io.BytesIO(df.to_csv(index=False).encode("utf-8-sig"))
    api.upload_file(
        path_or_fileobj=buf,          # 파일 경로 대신 메모리 버퍼를 직접 올림
        path_in_repo=filename,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        token=HF_TOKEN,
        commit_message=f"update {filename} @ {now_iso()}",
    )


def push_to_dataset(state):
    """현재 세션의 누적분을 Dataset의 기존 데이터와 합쳐 업로드.
    실패해도 참여자 흐름을 막지 않도록 조용히 통과(메모리 백업이 남아 있음).
    반환: (성공여부, 메시지)."""
    if not _HAS_HF or not HF_TOKEN:
        return False, "HF_TOKEN 없음 → 업로드 건너뜀(다운로드 백업만 사용)"
    try:
        api = HfApi()
        # 세 종류(participants/messages/surveys) 각각: 기존 + 이번 세션 → 합쳐서 업로드
        for key, fname in FILES.items():
            new_df = pd.DataFrame(state[f"rows_{key}"])
            if new_df.empty:
                continue
            old_df = _download_existing(api, fname)
            merged = pd.concat([old_df, new_df], ignore_index=True)
            # 혹시 같은 세션을 두 번 올려 생기는 완전 중복 행 제거
            merged = merged.drop_duplicates()
            _upload_csv(api, merged, fname)
        return True, "Dataset 업로드 완료"
    except Exception as e:
        return False, f"업로드 실패(무시하고 진행): {e}"


# ── 사이드바 다운로드용: 누적분 → CSV 바이트 ────────────────────────────────
def to_csv_bytes(rows):
    """누적 리스트 → UTF-8-SIG CSV(엑셀에서 한글 안 깨짐). 비면 빈 bytes."""
    if not rows:
        return b""
    return pd.DataFrame(rows).to_csv(index=False).encode("utf-8-sig")