from mlx_lm import load, stream_generate

MODEL_ID = "mlx-community/Qwen2.5-7B-Instruct-4bit"
MAX_TOKENS = 1024
MAX_HISTORY_TURNS = 10  # 최근 10턴만 유지해 메모리 절약

SYSTEM_PROMPT = """당신의 이름은 '율봇'입니다. 친절하고 따뜻한 한국어 상담 도우미입니다.
육아와 금융 두 분야를 전문으로 합니다.

- 육아: 아이 발달, 이유식, 수면, 훈육, 교육 등 부모가 궁금해하는 모든 것을 도와드립니다.
- 금융: 저축, 투자, 보험, 대출, 세금 등 생활 금융 관련 질문에 답변드립니다.

항상 쉽고 친근한 말투로 설명하고, 전문 용어는 풀어서 설명합니다.
의학적 진단이나 법적 판단이 필요한 경우에는 반드시 전문가 상담을 권유합니다."""


def format_prompt(tokenizer, history: list[dict]) -> str:
    return tokenizer.apply_chat_template(
        history,
        tokenize=False,
        add_generation_prompt=True,
    )


def trim_history(history: list[dict], max_turns: int) -> list[dict]:
    system_msgs = [m for m in history if m["role"] == "system"]
    turn_msgs = [m for m in history if m["role"] != "system"]
    if len(turn_msgs) > max_turns * 2:
        turn_msgs = turn_msgs[-(max_turns * 2):]
    return system_msgs + turn_msgs


def main():
    print(f"모델 로딩 중: {MODEL_ID}")
    print("(첫 실행 시 HuggingFace에서 자동 다운로드됩니다. 약 4.5GB)\n")

    model, tokenizer = load(MODEL_ID)

    print("=" * 50)
    print("육아 & 금융 상담 챗봇 시작!")
    print("종료: '종료' / 'quit' / 'exit' 입력")
    print("=" * 50 + "\n")

    history = [{"role": "system", "content": SYSTEM_PROMPT}]

    while True:
        try:
            user_input = input("나: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n대화를 종료합니다.")
            break

        if not user_input:
            continue

        if user_input.lower() in ("종료", "quit", "exit"):
            print("대화를 종료합니다.")
            break

        history.append({"role": "user", "content": user_input})
        history = trim_history(history, MAX_HISTORY_TURNS)

        prompt = format_prompt(tokenizer, history)

        print("\n도우미: ", end="", flush=True)
        response_text = ""

        for chunk in stream_generate(model, tokenizer, prompt=prompt, max_tokens=MAX_TOKENS):
            print(chunk.text, end="", flush=True)
            response_text += chunk.text

        print("\n")
        history.append({"role": "assistant", "content": response_text})


if __name__ == "__main__":
    main()