import openai
import time
import psutil
import json
import os

# GPU(VRAM) 사용량 측정을 위한 선택적 의존성 로드
try:
    import pynvml
    pynvml.nvmlInit()
    _GPU_AVAILABLE = True
except ModuleNotFoundError:
    _GPU_AVAILABLE = False

# -------------------------------------------
# (1) OCR 결과를 로드하여 인덱스-텍스트 딕셔너리 구축
# -------------------------------------------
# 실제 서비스에서는 OCR 엔진의 결과(JSON, TXT 등)를 로드하면 됩니다.
# 여기서는 예시로 하드코딩된 리스트를 사용합니다.

ocr_result = {
    "1": "2015年06月10日20:32",
    "2": "颜色分类：粘糖直径90厘米（有现货）",
    "3": "水车很漂亮",
    "4": "我们有专业的设计团队",
    "5": "旨在为亲们提供新颖的款式哦！",
    "6": "收起原图向左转向右转"
}

# 인덱스 부여 딕셔너리 생성
# ocr_dict = {idx: text for idx, text in enumerate(ocr_result, start=1)}
ocr_dict = ocr_result

# -------------------------------------------
# (2) LLM에게 전달할 프롬프트 작성
# -------------------------------------------


system_prompt = (
    "너는 중국어 상품 정보를 한국 쇼핑몰에 최적화된 형태로 번역하는 전문 번역가이자 카피라이터야. "
    "인덱스를 반드시 준수해야해"
    "아래에서 제공하는 `ocr_dict`의 값들은 모두 하나의 상품 이미지에서 추출한 중국어 문구(상품명·특징·설명 등)야. "
    "각 문구의 의미적 맥락을 고려해서 한국 소비자가 이해하기 쉬운 자연스러운 한국어로 변환해 줘. "
    "1) 상품명은 간결하고 매력적인 카피 문구로, 2) 특징·스펙은 한국어 어순·단위·표현에 맞춰, 3) 사용법·주의사항은 자연스러운 문장으로 번역해. "
    "중국어 브랜드/고유명사는 통용 표기(예: 小米→샤오미)를 사용하고, 숫자·단위를 대한민국 표준으로 바꿔. "
    "반드시 JSON만 반환해야 하며, 출력 형식은 예시를 따르되 주석·불필요한 문구를 포함하지 마. "
    "\n\n# 출력 예시\n{\n  \"1\": \"스마트 보온 컵받침대\",\n  \"2\": \"자동 감지 초고속 가열\",\n  \"3\": \"USB 전원, 휴대용 디자인\"\n}"
    "중국어 문장들을 한국어로 자연스럽고 의미가 잘 전달되게 번역해줘. "
)

# -------------------------------------------
# (3) OpenAI / vLLM 클라이언트 설정
# -------------------------------------------
host = "0.0.0.0"
port = "8180"
client = openai.Client(base_url=f"http://{host}:{port}/v1", api_key="null")

# -------------------------------------------
# (4) 메시지 구성
# -------------------------------------------
messages = [
    {"role": "system", "content": system_prompt},
    {
        "role": "user",
        "content": (
            "다음 중국어 문장들을 한국어로 자연스럽고 의미가 잘 전달되게 번역해줘. "
            "순서와 개수는 반드시 그대로 유지하고, 결과는 JSON 배열(리스트)로만 반환해. "
            "목적은 번역해줘. 결과에 중국어가 포함되어있으면 오답이야. 이럴경우 다시 번역해줘"
            "중국어 리스트:\n" + json.dumps(ocr_dict, ensure_ascii=False)
        ),
    },
]

# -------------------------------------------
# (5) 성능 측정용 헬퍼 함수들
# -------------------------------------------

def get_cpu_mem():
    """현재 프로세스의 CPU%, 메모리(RSS) MB"""
    process = psutil.Process(os.getpid())
    cpu = psutil.cpu_percent(interval=None)
    mem = process.memory_info().rss / (1024 ** 2)
    return cpu, mem


def get_gpu_mem():
    """첫 번째 GPU의 VRAM(MB) 사용량을 리턴, 미지원 시 None"""
    if not _GPU_AVAILABLE:
        return None
    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
    mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
    return mem_info.used / (1024 ** 2)


# -------------------------------------------
# (6) 모델 호출 + 트래킹
# -------------------------------------------

cpu_before, mem_before = get_cpu_mem()
vram_before = get_gpu_mem()
start_time = time.time()

# 스트리밍으로 응답 받으면서 출력 및 토큰 카운트 누적
response = client.chat.completions.create(
    model="null",
    messages=messages,
    stream=True,
)

output_tokens = 0

print("\n--- 번역 결과 (Streaming) ---")
for chunk in response:
    if chunk.choices[0].delta:
        content_piece = chunk.choices[0].delta.content
        output_tokens += len(content_piece.split())  # 단순 토큰 추정
        print(content_piece, end="")

print("\n-----------------------------\n")

# 후처리
elapsed = time.time() - start_time
cpu_after, mem_after = get_cpu_mem()
vram_after = get_gpu_mem()

input_tokens = sum(len(m["content"].split()) for m in messages)  # 단순 토큰 추정
throughput = output_tokens / elapsed if elapsed > 0 else 0

# -------------------------------------------
# (7) 메트릭 요약 출력
# -------------------------------------------

print("[메트릭 요약]")
print(f"• 응답 시간: {elapsed:.2f}s")
print(f"• CPU 사용량(전/후): {cpu_before:.1f}% → {cpu_after:.1f}%")
print(f"• 메인 메모리(전/후): {mem_before:.1f}MB → {mem_after:.1f}MB")
if vram_before is not None:
    print(f"• VRAM(전/후): {vram_before:.1f}MB → {vram_after:.1f}MB")
else:
    print("• VRAM: 측정 불가 (pynvml 미설치)")
print(f"• 입력 토큰: {input_tokens}개, 출력 토큰: {output_tokens}개")
print(f"• 처리량: {throughput:.2f} tokens/s")

# -------------------------------------------
# (8) 벤치마크 등 추가 처리 필요 시 여기에 구현
# -------------------------------------------