import google.generativeai as genai
import os

def classify_email(api_key, email_content, categories):
    """
    Classifies the given email content into one of the provided categories using the Gemini API.

    Args:
        api_key (str): The Gemini API key.
        email_content (str): The content of the email (subject + body).
        categories (list): A list of strings representing the possible categories.

    Returns:
        str: The chosen category, or None if classification fails.
    """
    if not api_key or api_key == "YOUR_GEMINI_API_KEY":
        print("! 경고: Gemini API 키가 설정되지 않았습니다. AI 분류를 건너뜁니다.")
        return None

    try:
        genai.configure(api_key=api_key)
    except Exception as e:
        print(f"Gemini API 설정 중 오류 발생: {e}")
        return None

    # For safety, let's use a model that's less likely to refuse classification
    model = genai.GenerativeModel('gemini-1.5-flash')

    # Construct the prompt
    prompt = f"""
    Analyze the following email content and classify it into one of the following categories.
    Your response must be ONLY ONE of the category names provided. Do not add any extra text, explanation, or punctuation.

    Categories: {', '.join(categories)}

    Email Content:
    ---
    {email_content}
    ---

    Category:
    """

    try:
        response = model.generate_content(prompt)
        # The response text should be one of the categories
        result_category = response.text.strip()
        
        # Validate that the model returned a valid category
        if result_category in categories:
            print(f"-> AI 분류 결과: '{result_category}'")
            return result_category
        else:
            print(f"! 경고: AI가 유효하지 않은 카테고리 '{result_category}'를 반환했습니다.")
            # Fallback or simply return None
            return None
            
    except Exception as e:
        print(f"Gemini API 호출 중 오류 발생: {e}")
        return None

if __name__ == '__main__':
    # --- 사용 예시 ---
    # 1. 아래 YOUR_API_KEY를 실제 키로 변경하세요.
    # 2. 터미널에서 `pip install google-generativeai` 를 실행하세요.
    # 3. `python ai_classifier.py` 명령으로 테스트하세요.
    
    test_api_key = os.environ.get("GEMINI_API_KEY", "YOUR_GEMINI_API_KEY") # 실제 키로 교체 필요
    
    test_email = """
    Subject: Invoice INV-2024-001 from Office Supplies Co.

    Dear Customer,

    Please find attached your invoice for the recent purchase of office supplies.
    The total amount due is $125.50.

    Thank you for your business.

    Sincerely,
    Office Supplies Co.
    """
    
    test_categories = ["청구서", "프로젝트 업데이트", "광고", "개인 용무"]
    
    if test_api_key != "YOUR_GEMINI_API_KEY":
        classified_category = classify_email(test_api_key, test_email, test_categories)
        
        if classified_category:
            print(f"\n최종 분류: {classified_category}")
        else:
            print("\n분류에 실패했습니다.")
    else:
        print("\n테스트를 위해 'YOUR_GEMINI_API_KEY'를 실제 API 키로 바꿔주세요.")