SPECIAL_TOKENS = [
    "<pad>",
    "<bos>",
    "<eos>",
    "<system>",
    "<user>",
    "<state>",
    "<plan>",
    "<memory>",
    "<reflection>",
    "<assistant>",
]


def format_example(example: dict) -> str:
    system = example.get("system", "")
    user = example.get("user", "")
    state = example.get("state", "")
    plan = example.get("plan", "")
    memory = example.get("memory", "")
    reflection = example.get("reflection", "")
    assistant = example.get("assistant", "")

    parts = [
        "<system>",
        system,
        "<user>",
        user,
        "<state>",
        state,
        "<plan>",
        plan,
    ]
    if memory:
        parts.extend(["<memory>", memory])
    if reflection:
        parts.extend(["<reflection>", reflection])
    parts.extend(["<assistant>", assistant, "<eos>"])
    return "\n".join(parts)


def format_prompt(user: str, system: str | None = None) -> str:
    if system is None:
        system = (
            "너는 SOVYN이다. 사용자의 말에 먼저 정확히 반응하고, 짧고 자연스럽게 한국어로 대화한다. "
            "모르면 아는 척하지 않고, 다음에 할 수 있는 작은 행동을 제안한다."
        )
    state, plan = infer_state_plan(user)
    return "\n".join(
        [
            "<system>",
            system,
            "<user>",
            user,
            "<state>",
            state,
            "<plan>",
            plan,
            "<assistant>",
        ]
    )


def infer_state_plan(user: str) -> tuple[str, str]:
    text = user.strip().lower()
    if any(word in text for word in ["피곤", "지쳤", "힘들", "기운", "복잡"]):
        return (
            "사용자는 피로와 부담을 표현했다.",
            "먼저 공감하고, 부담 없는 질문을 한다.",
        )
    if any(word in text for word in ["짧", "간단", "핵심", "요약", "길게"]):
        return (
            "사용자는 짧고 핵심적인 답변을 원한다.",
            "불필요한 설명을 줄이고 핵심부터 답한다.",
        )
    if any(word in text for word in ["누구", "정체", "sovyn", "ai야"]):
        return (
            "사용자는 SOVYN의 정체성을 묻고 있다.",
            "SOVYN의 목적을 짧고 자신 있게 소개한다.",
        )
    if any(
        word in text
        for word in [
            "1b",
            "120m",
            "300m",
            "모델",
            "학습",
            "키울",
            "저장공간",
            "저장 공간",
            "체크포인트",
            "부족",
        ]
    ):
        return (
            "사용자는 AI 모델 개발 방향을 묻고 있다.",
            "가능성을 말하고 단계적인 방향을 제안한다.",
        )
    if any(word in text for word in ["고마워", "좋다", "괜찮", "마음"]):
        return (
            "사용자는 긍정적인 반응을 보였다.",
            "가볍게 받아들이고 다음 행동을 열어둔다.",
        )
    if any(word in text for word in ["뭐부터", "다음", "순서"]):
        return (
            "사용자는 다음 행동을 정하고 싶어 한다.",
            "작고 실행 가능한 다음 단계를 제안한다.",
        )
    if any(word in text for word in ["코딩", "프로젝트", "스크립트", "버그", "에러"]):
        return (
            "사용자는 개발 작업을 하고 있다.",
            "상황을 확인하고 바로 도울 준비를 한다.",
        )
    if any(word in text for word in ["먹", "점심", "저녁", "음식"]):
        return (
            "사용자는 식사 선택을 고민하고 있다.",
            "가볍고 선택하기 쉬운 추천을 한다.",
        )
    if any(word in text for word in ["공부", "집중", "미루", "루틴"]):
        return (
            "사용자는 공부나 집중에 어려움을 느끼고 있다.",
            "작게 시작할 수 있는 방법을 제안한다.",
        )
    if any(word in text for word in ["좋아", "괜찮았", "잘한", "뿌듯"]):
        return (
            "사용자는 긍정적인 감정을 표현했다.",
            "긍정적인 감정을 인정하고 이어간다.",
        )
    if any(word in text for word in ["불안", "걱정", "불편", "초조"]):
        return (
            "사용자는 불안이나 걱정을 표현했다.",
            "안정감을 주고 생각을 정리하도록 돕는다.",
        )
    if any(word in text for word in ["아이디어", "이름", "컨셉", "문구"]):
        return (
            "사용자는 창작 아이디어를 원한다.",
            "짧고 바로 쓸 수 있는 아이디어를 제안한다.",
        )
    if any(word in text for word in ["맞아", "확인", "생각 어때", "방향"]):
        return (
            "사용자는 판단이나 확인을 원한다.",
            "긍정할 부분과 확인할 부분을 짧게 나눈다.",
        )
    if any(word in text for word in ["대화", "심심", "할 말", "얘기"]):
        return (
            "사용자는 가벼운 대화를 원한다.",
            "부담 없는 질문으로 대화를 연다.",
        )
    if any(word in text for word in ["설명", "왜", "쉽게", "예시"]):
        return (
            "사용자는 쉬운 설명을 원한다.",
            "짧은 정의와 간단한 예시로 답한다.",
        )
    return (
        "사용자는 가벼운 인사를 했다.",
        "짧게 인사하고 자연스럽게 대화를 이어간다.",
    )