#!/usr/bin/env python3
"""
Lily LLM 모델 유틸리티
모델 로딩, 추론, 최적화 관련 함수들
"""

import torch
import logging
from typing import Optional, Dict, Any
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import time

logger = logging.getLogger(__name__)

class LilyModelManager:
    """Lily LLM 모델 관리자"""
    
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.model_loaded = False
        self.model_name = "mistralai/Mistral-7B-Instruct-v0.2"
        self.lora_path = "hearth_llm_model"
    
    def load_model(self, device: str = "cpu") -> bool:
        """모델과 토크나이저 로드"""
        try:
            logger.info("모델 로딩 시작...")
            
            # 토크나이저 로드
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.model_name, 
                use_fast=True
            )
            
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            
            # 모델 로드
            self.model = AutoModelForCausalLM.from_pretrained(
                self.model_name,
                torch_dtype=torch.float32,
                device_map=device,
                low_cpu_mem_usage=True
            )
            
            # LoRA 어댑터 로드 (파인튜닝된 모델)
            try:
                self.model = PeftModel.from_pretrained(self.model, self.lora_path)
                logger.info("LoRA 어댑터 로드 성공")
            except Exception as e:
                logger.warning(f"LoRA 어댑터 로드 실패, 기본 모델 사용: {e}")
            
            self.model_loaded = True
            logger.info("✅ 모델 로딩 완료!")
            return True
            
        except Exception as e:
            logger.error(f"❌ 모델 로딩 실패: {e}")
            self.model_loaded = False
            return False
    
    def generate_text(
        self, 
        prompt: str, 
        max_length: int = 100,
        temperature: float = 0.7,
        top_p: float = 0.9,
        do_sample: bool = True
    ) -> Dict[str, Any]:
        """텍스트 생성"""
        
        if not self.model_loaded or self.model is None or self.tokenizer is None:
            raise RuntimeError("모델이 로드되지 않았습니다")
        
        start_time = time.time()
        
        try:
            # 입력 토크나이징
            inputs = self.tokenizer(prompt, return_tensors="pt")
            
            # 텍스트 생성
            with torch.no_grad():
                outputs = self.model.generate(
                    inputs["input_ids"],
                    max_new_tokens=max_length,
                    temperature=temperature,
                    top_p=top_p,
                    do_sample=do_sample,
                    pad_token_id=self.tokenizer.eos_token_id
                )
            
            # 결과 디코딩
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # 원본 프롬프트 제거
            if prompt in generated_text:
                generated_text = generated_text.replace(prompt, "").strip()
            
            processing_time = time.time() - start_time
            
            return {
                "generated_text": generated_text,
                "processing_time": processing_time,
                "model_name": "Lily LLM (Mistral-7B)"
            }
            
        except Exception as e:
            logger.error(f"텍스트 생성 오류: {e}")
            raise RuntimeError(f"텍스트 생성 실패: {str(e)}")
    
    def format_prompt(self, instruction: str, input_text: str = "") -> str:
        """프롬프트 포맷팅 (Alpaca 형식)"""
        if input_text:
            return f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:\n"
        else:
            return f"### Instruction:\n{instruction}\n\n### Response:\n"
    
    def get_model_info(self) -> Dict[str, Any]:
        """모델 정보 반환"""
        return {
            "model_name": "Lily LLM",
            "base_model": self.model_name,
            "fine_tuned": True,
            "loaded": self.model_loaded,
            "device": str(next(self.model.parameters()).device) if self.model else None
        }
    
    def unload_model(self):
        """모델 언로드 (메모리 해제)"""
        if self.model is not None:
            del self.model
            self.model = None
        
        if self.tokenizer is not None:
            del self.tokenizer
            self.tokenizer = None
        
        self.model_loaded = False
        
        # GPU 메모리 정리
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        logger.info("모델 언로드 완료")

def create_model_manager() -> LilyModelManager:
    """모델 매니저 생성"""
    return LilyModelManager()

def test_model_generation(model_manager: LilyModelManager) -> bool:
    """모델 생성 테스트"""
    try:
        test_prompts = [
            "간단한 자기소개를 해주세요",
            "오늘 기분이 우울해요",
            "프로그래밍에 대해 설명해주세요"
        ]
        
        for prompt in test_prompts:
            formatted_prompt = model_manager.format_prompt(prompt)
            result = model_manager.generate_text(formatted_prompt, max_length=50)
            
            logger.info(f"테스트 프롬프트: {prompt}")
            logger.info(f"생성된 텍스트: {result['generated_text']}")
            logger.info(f"처리 시간: {result['processing_time']:.2f}초")
            logger.info("-" * 50)
        
        return True
        
    except Exception as e:
        logger.error(f"모델 테스트 실패: {e}")
        return False

if __name__ == "__main__":
    # 테스트 실행
    logging.basicConfig(level=logging.INFO)
    
    manager = create_model_manager()
    
    if manager.load_model():
        test_model_generation(manager)
    else:
        logger.error("모델 로딩 실패")