"""modal_gpu_ez 출력을 LangChain 형식으로 변환하는 유틸리티."""

from __future__ import annotations

from typing import Any


def strip_prompt_prefix(prompt: str, generated_text: str) -> str:
    """HF text-generation 파이프라인은 프롬프트를 에코하므로 제거한다."""
    if generated_text.startswith(prompt):
        return generated_text[len(prompt) :]
    return generated_text


def apply_stop_words(text: str, stop: list[str] | None) -> str:
    """첫 번째 stop word 위치에서 텍스트를 절단한다."""
    if not stop:
        return text
    min_idx = len(text)
    for s in stop:
        idx = text.find(s)
        if idx != -1 and idx < min_idx:
            min_idx = idx
    return text[:min_idx]


def parse_text_generation_output(
    raw: Any,
    prompt: str,
    stop: list[str] | None,
) -> str:
    """HF 파이프라인 출력을 파싱하여 순수 생성 텍스트를 반환한다."""
    # text-generation: [{"generated_text": "..."}]
    if isinstance(raw, list) and raw and isinstance(raw[0], dict):
        text = raw[0].get("generated_text", "")
    # translation: ["translated text"]
    elif isinstance(raw, list) and raw and isinstance(raw[0], str):
        text = raw[0]
    else:
        text = str(raw)

    text = strip_prompt_prefix(prompt, text)
    text = apply_stop_words(text, stop)
    return text
