import os
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
os.environ["HF_HOME"] = "/tmp/hf_cache"
os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
import json
import requests
import streamlit as st
from datetime import datetime
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline



# ✅ Streamlit 기본 설정
st.set_page_config(page_title="학사일정 캘린더", layout="centered")
st.title("📅 학사일정 캘린더 + AI 요약")
st.markdown("NEIS API에서 학사일정을 불러오고 FullCalendar로 시각화합니다.")

# ✅ 디버깅용 토큰 확인
st.write("✅ 캐시 경로:", os.environ.get("TRANSFORMERS_CACHE"))
st.write("🔐 토큰 있음:", os.environ.get("HUGGINGFACE_TOKEN") is not None)

# ✅ Gemma 모델 로딩 함수
@st.cache_resource
def load_model():
    token = os.environ.get("HUGGINGFACE_TOKEN")
    model_id = "google/gemma-2-2b-it"
    cache_dir = "/tmp/hf_cache"  # ✅ Hugging Face Spaces에서 허용된 디렉토리

    tokenizer = AutoTokenizer.from_pretrained(
        model_id,
        use_auth_token=token,
        cache_dir=cache_dir
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        use_auth_token=token,
        cache_dir=cache_dir
    )

    return pipeline("text-generation", model=model, tokenizer=tokenizer)

llm = load_model()

# ✅ 학교 정보 가져오기
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

# ✅ AI 요약 함수
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()

# ✅ UI: 지역 / 학교 / 연도 선택
region_options = {
    "B10": "서울", "C10": "부산", "D10": "대구", "E10": "인천", "F10": "광주", "G10": "대전",
    "H10": "울산", "I10": "세종", "J10": "경기", "K10": "강원", "M10": "충북", "N10": "충남",
    "P10": "전북", "Q10": "전남", "R10": "경북", "S10": "경남", "T10": "제주"
}

with st.form("query_form"):
    region = st.selectbox("지역 교육청", options=list(region_options.keys()), format_func=lambda x: f"{region_options[x]} ({x})")
    school_name = st.text_input("학교명", placeholder="예: 상리초등학교")
    year = st.selectbox("년도", options=[2022, 2023, 2024, 2025], index=2)
    submitted = st.form_submit_button("📅 학사일정 불러오기")

# ✅ 제출 시 처리
if submitted:
    with st.spinner("일정 불러오는 중..."):
        api_key = os.environ.get("NEIS_API_KEY", "a69e08342c8947b4a52cd72789a5ecaf")
        school_code, region_code = get_school_info(region, school_name, api_key)
        if not school_code:
            st.error("학교 정보를 찾을 수 없습니다.")
        else:
            schedule_rows = get_schedule(region_code, school_code, year, api_key)
            if not schedule_rows:
                st.info("해당 조건의 학사일정이 없습니다.")
            else:
                # 캘린더용 JSON 생성
                events = [
                    {
                        "title": row["EVENT_NM"],
                        "start": datetime.strptime(row["AA_YMD"], "%Y%m%d").strftime("%Y-%m-%d")
                    }
                    for row in schedule_rows
                ]
                event_json = json.dumps(events, ensure_ascii=False)

                # FullCalendar 삽입
                st.components.v1.html(f"""
                <html>
                <head>
                  <link href='https://cdn.jsdelivr.net/npm/fullcalendar@6.1.8/index.global.min.css' rel='stylesheet' />
                  <script src='https://cdn.jsdelivr.net/npm/fullcalendar@6.1.8/index.global.min.js'></script>
                  <script>
                    document.addEventListener('DOMContentLoaded', function() {{
                      var calendarEl = document.getElementById('calendar');
                      var calendar = new FullCalendar.Calendar(calendarEl, {{
                        initialView: 'dayGridMonth',
                        locale: 'ko',
                        height: 600,
                        events: {event_json}
                      }});
                      calendar.render();
                    }});
                  </script>
                </head>
                <body>
                  <div id='calendar'></div>
                </body>
                </html>
                """, height=650)

                # 요약 버튼
                with st.expander("✨ 1년치 요약 보기", expanded=False):
                    if st.button("🤖 요약 생성하기"):
                        with st.spinner("Gemma 모델이 요약 중..."):
                            summary = summarize_schedule(schedule_rows, school_name, year)
                            st.success("요약 완료!")
                            st.markdown(f"**{school_name} {year}년 학사일정 요약:**\n\n{summary}")