app.py

import os
import requests
import gradio as gr
from datetime import datetime
from transformers import pipeline

# ────────────────────────────── 모델 불러오기 ──────────────────────────────
llm = pipeline("text-generation", model="google/gemma-2-2b-it", max_new_tokens=512)

# ────────────────────────────── API 처리 함수 ──────────────────────────────
def get_school_info(region_code, school_name, api_key):
    url = f"https://open.neis.go.kr/hub/schoolInfo?KEY={api_key}&Type=json&pIndex=1&pSize=1&SCHUL_NM={school_name}&ATPT_OFCDC_SC_CODE={region_code}"
    res = requests.get(url)
    data = res.json()
    school = data.get("schoolInfo", [{}])[1].get("row", [{}])[0]
    return school.get("SD_SCHUL_CODE"), school.get("ATPT_OFCDC_SC_CODE")

def get_schedule(region_code, school_code, year, api_key):
    from_ymd = f"{year}0101"
    to_ymd = f"{year}1231"
    url = f"https://open.neis.go.kr/hub/SchoolSchedule?KEY={api_key}&Type=json&pIndex=1&pSize=500&ATPT_OFCDC_SC_CODE={region_code}&SD_SCHUL_CODE={school_code}&AA_FROM_YMD={from_ymd}&AA_TO_YMD={to_ymd}"
    res = requests.get(url)
    data = res.json()
    rows = data.get("SchoolSchedule", [{}])[1].get("row", [])
    return rows

def summarize_schedule(rows, school_name, year):
    lines = []
    for row in rows:
        date = row["AA_YMD"]
        dt = datetime.strptime(date, "%Y%m%d").strftime("%-m월 %-d일")
        event = row["EVENT_NM"]
        lines.append(f"{dt}: {event}")
    text = "\n".join(lines)
    prompt = f"{school_name}의 {year}년 학사일정은 다음과 같다:\n{text}\n이 내용을 요약해줘. 주요 일정 중심으로 간결하게."
    result = llm([{"role": "user", "content": prompt}])
    return result[0]["generated_text"].replace(prompt, "").strip()

# ────────────────────────────── 그라디오 인터페이스 함수 ──────────────────────────────
def load_and_summarize(region_code, school_name, year, month):
    api_key = os.environ.get("NEIS_API_KEY", "a69e08342c8947b4a52cd72789a5ecaf")

    if not school_name.strip():
        return "❗ 학교명을 입력해주세요.", "", ""

    school_code, confirmed_region = get_school_info(region_code, school_name, api_key)
    if not school_code:
        return "❗ 학교 정보를 찾을 수 없습니다. 입력을 확인해주세요.", "", ""

    schedule_rows = get_schedule(confirmed_region, school_code, year, api_key)
    if not schedule_rows:
        return "ℹ️ 학사일정 정보가 없습니다.", "", ""

    # 월별 일정 필터링
    month_rows = [r for r in schedule_rows if r["AA_YMD"].startswith(f"{year}{month:02}")]
    month_display = "\n".join(
        f"{datetime.strptime(r['AA_YMD'], '%Y%m%d').strftime('%Y-%m-%d')} - {r['EVENT_NM']}"
        for r in month_rows
    ) or "📭 선택한 월에 해당하는 일정이 없습니다."

    # 요약 생성
    summary = summarize_schedule(schedule_rows, school_name, year)

    return f"✅ {school_name}의 {year}년 {month}월 일정입니다.", month_display, summary

# ────────────────────────────── Gradio UI 구성 ──────────────────────────────
region_options = {
    "B10": "서울", "C10": "부산", "D10": "대구", "E10": "인천", "F10": "광주", "G10": "대전",
    "H10": "울산", "I10": "세종", "J10": "경기", "K10": "강원", "M10": "충북", "N10": "충남",
    "P10": "전북", "Q10": "전남", "R10": "경북", "S10": "경남", "T10": "제주"
}

with gr.Blocks(title="학사일정 요약 캘린더") as demo:
    gr.Markdown("## 🏫 학사일정 조회 및 AI 요약 (google/gemma-2-2b-it)")
    with gr.Row():
        region = gr.Dropdown(label="지역 교육청", choices=list(region_options.keys()), value="B10", interactive=True)
        school_name = gr.Textbox(label="학교명", placeholder="예: 상리초등학교")
    with gr.Row():
        year = gr.Dropdown(label="년도", choices=[2022, 2023, 2024, 2025], value=2024)
        month = gr.Slider(label="월", minimum=1, maximum=12, value=7, step=1)

    btn = gr.Button("📅 일정 불러오기 및 요약")

    status = gr.Markdown()
    month_output = gr.Textbox(label="📆 선택한 월의 일정", lines=8)
    summary_output = gr.Textbox(label="📝 1년 학사일정 요약 (Gemma)", lines=6)

    btn.click(fn=load_and_summarize, inputs=[region, school_name, year, month], outputs=[status, month_output, summary_output])

demo.launch()