import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import time
import logging
import os
import json
from datetime import datetime

# --- 설정 ---
# ★★★ 모델 가중치와 토크나이저 ID 분리 ★★★
MODEL_ID_FOR_WEIGHTS = "unsloth/gemma-3-1b-it-bnb-4bit" # 모델 가중치는 여기서 로드
TOKENIZER_ID = "google/gemma-3-1b-it"               # 토크나이저는 원본 여기서 로드

# CPU 사용 (HF Spaces 무료 티어 기준)
DEVICE = "cpu"
# 메모리 파일 경로
MEMORY_FILE = "thought_memory.json"
# 생각 주기 (초)
THINKING_INTERVAL_SECONDS = 120 # 예: 2분마다 생각
# 생성할 최대 토큰 수
MAX_NEW_TOKENS = 150
# 초기 생각 프롬프트
INITIAL_PROMPT = "나는 계속해서 스스로 생각하는 AI입니다. 나의 첫 번째 생각은 다음과 같습니다:"

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# (load_memory, save_memory 함수는 이전과 동일)
def load_memory():
    """메모리 파일에서 이전 생각 기록을 로드합니다."""
    if os.path.exists(MEMORY_FILE):
        try:
            with open(MEMORY_FILE, 'r', encoding='utf-8') as f:
                memory = json.load(f)
                if not isinstance(memory, list):
                    logging.warning(f"{MEMORY_FILE} 내용이 리스트가 아니므로 초기화합니다.")
                    return []
                logging.info(f"{len(memory)}개의 이전 생각을 로드했습니다.")
                return memory
        except json.JSONDecodeError:
            logging.error(f"{MEMORY_FILE} 파일 파싱 오류. 메모리를 초기화합니다.")
            return []
        except Exception as e:
            logging.error(f"메모리 로드 중 오류 발생: {e}", exc_info=True)
            return []
    else:
        logging.info("메모리 파일이 없어 새로 시작합니다.")
        return []

def save_memory(memory):
    """현재 생각 기록을 메모리 파일에 저장합니다."""
    try:
        with open(MEMORY_FILE, 'w', encoding='utf-8') as f:
            json.dump(memory, f, ensure_ascii=False, indent=2)
        logging.debug(f"메모리를 {MEMORY_FILE}에 저장했습니다.")
    except Exception as e:
        logging.error(f"메모리 저장 중 오류 발생: {e}", exc_info=True)


def generate_thought(tokenizer, model, prompt_history):
    """주어진 프롬프트 기록을 바탕으로 다음 생각을 생성합니다."""
    if not prompt_history:
        chat = [{"role": "user", "content": INITIAL_PROMPT}]
    else:
        last_thought = prompt_history[-1]['content']
        prompt = f"이전 생각: \"{last_thought}\"\n\n이 생각을 바탕으로 다음으로 떠오르는 생각이나 질문, 또는 확장된 개념은 무엇인가요? 간결하게 답해주세요."
        chat = [{"role": "user", "content": prompt}]

    prompt_formatted = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
    logging.info(f"--- 모델 입력 프롬프트 ---\n{prompt_formatted}\n-----------------------")

    inputs = tokenizer(prompt_formatted, return_tensors="pt").to(DEVICE)

    start_time = time.time()
    logging.info("모델 추론 시작...")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            pad_token_id=tokenizer.eos_token_id
        )
    end_time = time.time()
    logging.info(f"모델 추론 완료 ({end_time - start_time:.2f}초 소요)")

    input_token_length = inputs.input_ids.shape[1]
    generated_ids = outputs[0, input_token_length:]
    new_thought_raw = tokenizer.decode(generated_ids, skip_special_tokens=True)

    logging.info(f"모델 생성 결과 (Raw): {new_thought_raw}")
    return new_thought_raw.strip()

if __name__ == "__main__":
    logging.info("AI 생각 프로세스 시작...")
    logging.info(f"Tokenizer ID: {TOKENIZER_ID}")
    logging.info(f"Model Weights ID: {MODEL_ID_FOR_WEIGHTS}")
    logging.info(f"실행 장치: {DEVICE}")

    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        logging.info("Hugging Face 토큰을 사용합니다.")
    else:
        logging.info("Hugging Face 토큰이 설정되지 않았습니다 (필요 시 Secrets에 추가).")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 # CPU 지원 안되면 float32로 변경
        # bnb_4bit_compute_dtype=torch.float32 # bfloat16 문제 시 이 라인 사용
    )

    try:
        logging.info("토크나이저 로딩 중...")
        # ★★★ 토크나이저 로딩 시 원본 ID 사용 ★★★
        tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_ID, token=hf_token)

        logging.info("양자화된 모델 로딩 중... (bitsandbytes 설정 적용)")
        # ★★★ 모델 가중치 로딩 시 양자화 모델 ID 사용 ★★★
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID_FOR_WEIGHTS,
            quantization_config=bnb_config,
            device_map=DEVICE,
            token=hf_token
        )

        model.eval()
        logging.info("모델 및 토크나이저 로드 완료.")

    except Exception as e:
        logging.error(f"모델 또는 토크나이저 로딩 중 치명적 오류 발생: {e}", exc_info=True)
        # CPU가 bfloat16 지원하지 않으면 여기서 오류 발생 가능
        exit(1)

    thought_history = load_memory()

    try:
        while True:
            logging.info("=== 새로운 생각 사이클 시작 ===")
            new_thought = generate_thought(tokenizer, model, thought_history)

            if new_thought:
                logging.info(f"생성된 새로운 생각: {new_thought}")
                thought_entry = {"role": "assistant", "content": new_thought, "timestamp": datetime.now().isoformat()}
                thought_history.append(thought_entry)
                save_memory(thought_history)
            else:
                logging.warning("모델이 빈 생각을 생성했습니다.")

            logging.info(f"다음 생각까지 {THINKING_INTERVAL_SECONDS}초 대기...")
            time.sleep(THINKING_INTERVAL_SECONDS)

    except KeyboardInterrupt:
        logging.info("사용자 요청으로 AI 프로세스 중지.")
    except Exception as e:
        logging.error(f"메인 루프에서 오류 발생: {e}", exc_info=True)
    finally:
        logging.info("AI 생각 프로세스 종료. 최종 메모리 저장 시도.")
        save_memory(thought_history)