"""
EXAONE-3.0-7.8B-Instruct 모델을 Ollama에 추가하는 스크립트 (토큰 포함)
"""

import os
import subprocess
from pathlib import Path
import sys

# 프로젝트 루트를 Python 경로에 추가
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def get_huggingface_token():
    """
    Hugging Face 토큰 가져오기 (환경 변수 우선, 없으면 DB에서)
    Google API 키 관리 방식과 동일한 패턴
    """
    # 환경 변수에서 먼저 확인
    token = os.getenv('HUGGINGFACE_HUB_TOKEN', '').strip()
    if not token:
        token = os.getenv('HF_TOKEN', '').strip()
    
    if token:
        print(f"[Hugging Face] 환경 변수에서 토큰 가져옴 (길이: {len(token)}자)")
        return token
    
    # DB에서 가져오기 (순환 참조 방지를 위해 여기서 임포트)
    try:
        from app.database import SystemConfig
        # Flask 애플리케이션 컨텍스트 없이도 작동하도록 시도
        token = SystemConfig.get_config('huggingface_token', '').strip()
        if token:
            print(f"[Hugging Face] DB에서 토큰 가져옴 (길이: {len(token)}자)")
        return token
    except Exception as e:
        print(f"[Hugging Face] DB에서 토큰 조회 실패 (환경 변수 사용): {e}")
        return ''

# Hugging Face 토큰 가져오기
HF_TOKEN = get_huggingface_token()
if not HF_TOKEN:
    raise ValueError("HUGGINGFACE_HUB_TOKEN 또는 HF_TOKEN 환경 변수를 설정하거나, 관리 페이지에서 토큰을 설정해주세요.")

def set_huggingface_token():
    """Hugging Face 토큰 설정"""
    os.environ['HUGGINGFACE_HUB_TOKEN'] = HF_TOKEN
    print(f"[OK] Hugging Face 토큰 설정 완료")

def create_ollama_modelfile():
    """Ollama Modelfile 생성"""
    modelfile_content = f"""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"[OK] Modelfile 생성 완료: {modelfile_path.absolute()}")
    return modelfile_path

def create_model_with_modelfile(modelfile_path):
    """Modelfile을 사용하여 Ollama 모델 생성"""
    try:
        print("\n" + "=" * 60)
        print("Ollama 모델 생성 시작...")
        print("=" * 60)
        
        # 환경 변수 설정
        env = os.environ.copy()
        env['HUGGINGFACE_HUB_TOKEN'] = HF_TOKEN
        
        # ollama create 명령어 실행
        cmd = ['ollama', 'create', 'EXAONE-3.0-7.8B-Instruct', '-f', str(modelfile_path)]
        print(f"실행 명령어: {' '.join(cmd)}")
        
        result = subprocess.run(
            cmd,
            env=env,
            capture_output=True,
            text=True,
            timeout=3600  # 1시간 타임아웃
        )
        
        if result.returncode == 0:
            print("[OK] 모델 생성 성공!")
            print(result.stdout)
            return True
        else:
            print("[ERROR] 모델 생성 실패")
            print("오류 출력:")
            print(result.stderr)
            return False
            
    except subprocess.TimeoutExpired:
        print("[ERROR] 모델 생성 시간 초과 (1시간)")
        return False
    except Exception as e:
        print(f"[ERROR] 오류 발생: {e}")
        return False

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

def verify_model():
    """생성된 모델 확인"""
    try:
        result = subprocess.run(['ollama', 'list'], 
                              capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            if 'EXAONE-3.0-7.8B-Instruct' in result.stdout:
                print("\n[OK] 모델이 성공적으로 추가되었습니다!")
                print("\n설치된 모델 목록:")
                print(result.stdout)
                return True
            else:
                print("\n[WARNING] 모델이 목록에 나타나지 않습니다.")
                print("\n현재 설치된 모델:")
                print(result.stdout)
                return False
        return False
    except Exception as e:
        print(f"[WARNING] 모델 확인 중 오류: {e}")
        return False

if __name__ == "__main__":
    print("\n" + "=" * 60)
    print("EXAONE-3.0-7.8B-Instruct Ollama 추가 스크립트")
    print("=" * 60 + "\n")
    
    # 사전 확인
    if not check_ollama_installation():
        print("\n[WARNING] Ollama를 먼저 설치해주세요.")
        exit(1)
    
    # 토큰 설정
    set_huggingface_token()
    
    # Modelfile 생성
    modelfile_path = create_ollama_modelfile()
    
    # 모델 생성 확인
    print("\n[WARNING] 주의사항:")
    print("- 모델 크기가 약 15GB이므로 다운로드에 시간이 걸릴 수 있습니다.")
    print("- 충분한 디스크 공간과 GPU 메모리가 필요합니다.")
    print("- Hugging Face에서 모델 액세스 권한이 필요합니다.")
    
    # 모델 생성 시작
    print("\n모델 생성을 시작합니다...")
    
    # 모델 생성
    success = create_model_with_modelfile(modelfile_path)
    
    if success:
        verify_model()
        print("\n" + "=" * 60)
        print("[OK] 모델 추가 완료!")
        print("=" * 60)
        print("\n이제 웹 애플리케이션에서 모델을 사용할 수 있습니다.")
        print("모델 테스트:")
        print("  ollama run EXAONE-3.0-7.8B-Instruct \"안녕하세요\"")
        print("=" * 60)
    else:
        print("\n" + "=" * 60)
        print("[ERROR] 모델 생성 실패")
        print("=" * 60)
        print("\n수동으로 모델을 생성하려면:")
        print(f"  1. Hugging Face 토큰 설정:")
        print(f"     $env:HUGGINGFACE_HUB_TOKEN='YOUR_TOKEN_HERE'")
        print(f"  2. Modelfile로 모델 생성:")
        print(f"     ollama create EXAONE-3.0-7.8B-Instruct -f {modelfile_path}")
        print("=" * 60)