"""
EXAONE-3.0-7.8B-Instruct 모델을 Ollama에 추가하는 스크립트

사용 방법:
1. Hugging Face에서 모델 액세스 권한을 받아야 합니다.
2. 필요한 패키지를 설치합니다:
   pip install transformers torch huggingface_hub
3. 이 스크립트를 실행합니다:
   python add_exaone_model.py
"""

import os
import subprocess
import json
from pathlib import Path

def create_ollama_modelfile():
    """Ollama Modelfile 생성"""
    # EXAONE 모델의 실제 채팅 템플릿 형식에 맞춰 수정
    modelfile_content = """FROM huggingface:LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct

# 모델 설정
PARAMETER temperature 0.7
PARAMETER top_p 0.9
PARAMETER top_k 40
PARAMETER num_ctx 4096
PARAMETER num_predict 512

# 시스템 프롬프트
SYSTEM \"\"\"You are EXAONE, a helpful AI assistant developed by LG AI Research. 
You are designed to be helpful, harmless, and honest. 
You can communicate in both Korean and English.\"\"\"

# EXAONE 모델의 채팅 템플릿 형식
# 참고: 실제 모델의 토크나이저 템플릿을 확인하여 조정이 필요할 수 있습니다.
TEMPLATE \"\"\"{{ if .System }}<|im_start|>system
{{ .System }}<|im_end|>
{{ end }}{{ if .Prompt }}<|im_start|>user
{{ .Prompt }}<|im_end|>
{{ end }}<|im_start|>assistant
{{ .Response }}<|im_end|>
\"\"\"
"""
    
    modelfile_path = Path("EXAONE-3.0-7.8B-Instruct.modelfile")
    modelfile_path.write_text(modelfile_content, encoding='utf-8')
    print(f"✅ Modelfile 생성 완료: {modelfile_path}")
    return modelfile_path

def create_ollama_model_from_huggingface():
    """Hugging Face 모델을 사용하여 Ollama 모델 생성"""
    print("=" * 60)
    print("EXAONE-3.0-7.8B-Instruct 모델을 Ollama에 추가합니다")
    print("=" * 60)
    
    # 방법 1: Modelfile 사용 (권장)
    print("\n[방법 1] Modelfile을 사용한 모델 생성")
    print("-" * 60)
    
    modelfile_path = create_ollama_modelfile()
    
    print(f"\n다음 명령어를 실행하여 모델을 생성하세요:")
    print(f"  ollama create EXAONE-3.0-7.8B-Instruct -f {modelfile_path}")
    
    # 방법 2: 직접 Hugging Face에서 가져오기
    print("\n[방법 2] Hugging Face에서 직접 가져오기")
    print("-" * 60)
    print("다음 명령어를 실행하세요:")
    print("  ollama create EXAONE-3.0-7.8B-Instruct --from huggingface:LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct")
    
    # 방법 3: Python 스크립트로 변환
    print("\n[방법 3] GGUF 형식으로 변환 후 추가 (고급)")
    print("-" * 60)
    print("이 방법은 llama.cpp를 사용하여 모델을 GGUF 형식으로 변환한 후")
    print("Ollama에 추가하는 방법입니다. 더 많은 설정이 필요합니다.")
    
    print("\n" + "=" * 60)
    print("⚠️  중요 사항:")
    print("=" * 60)
    print("1. Hugging Face에서 모델 액세스 권한이 필요합니다.")
    print("   https://huggingface.co/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct")
    print("2. Hugging Face 토큰이 필요할 수 있습니다.")
    print("   환경 변수에 설정: export HUGGINGFACE_HUB_TOKEN=your_token")
    print("   Windows PowerShell: $env:HUGGINGFACE_HUB_TOKEN='your_token'")
    print("3. 모델 크기가 약 15GB이므로 충분한 디스크 공간이 필요합니다.")
    print("4. GPU 메모리가 충분해야 합니다 (최소 16GB 권장).")
    print("=" * 60)

def check_ollama_installation():
    """Ollama 설치 확인"""
    try:
        result = subprocess.run(['ollama', '--version'], 
                              capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            print(f"✅ Ollama 설치 확인: {result.stdout.strip()}")
            return True
        else:
            print("❌ Ollama가 설치되어 있지 않습니다.")
            return False
    except FileNotFoundError:
        print("❌ Ollama가 설치되어 있지 않습니다.")
        print("   설치 방법: https://ollama.ai/download")
        return False
    except Exception as e:
        print(f"❌ Ollama 확인 중 오류: {e}")
        return False

def check_huggingface_access():
    """Hugging Face 액세스 확인"""
    try:
        from huggingface_hub import whoami
        user_info = whoami()
        print(f"✅ Hugging Face 로그인 확인: {user_info.get('name', 'Unknown')}")
        return True
    except Exception as e:
        print(f"⚠️  Hugging Face 로그인 확인 실패: {e}")
        print("   Hugging Face CLI로 로그인하세요: huggingface-cli login")
        return False

if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("EXAONE-3.0-7.8B-Instruct Ollama 추가 스크립트")
    print("=" * 60 + "\n")
    
    # 사전 확인
    ollama_ok = check_ollama_installation()
    hf_ok = check_huggingface_access()
    
    if not ollama_ok:
        print("\n⚠️  Ollama를 먼저 설치해주세요.")
        exit(1)
    
    # 모델 추가 방법 안내
    create_ollama_model_from_huggingface()
    
    print("\n" + "=" * 60)
    print("다음 단계:")
    print("=" * 60)
    print("1. 위의 명령어 중 하나를 선택하여 실행하세요.")
    print("2. 모델 생성이 완료되면 다음 명령어로 확인하세요:")
    print("   ollama list")
    print("3. 모델이 목록에 나타나면 웹 애플리케이션에서 사용할 수 있습니다.")
    print("=" * 60 + "\n")