import os
import subprocess
import requests
import gradio as gr
from datetime import datetime
from llama_cpp import Llama  # GPT4ALL용

# ───────────────────────── 모델 자동 다운로드 ─────────────────────────
MODEL_PATH = "ggjt-model.bin"
GOOGLE_DRIVE_FILE_ID = "1OTd4JOU2ZYC9fU6BP9HZb_8OhCFMq8Ch"

if not os.path.exists(MODEL_PATH):
    print("⏬ ggjt-model.bin 다운로드 중...")
    subprocess.run(["gdown", "--id", GOOGLE_DRIVE_FILE_ID, "-O", MODEL_PATH])
    print("✅ 다운로드 완료.")

# ───────────────────────── 모델 로딩 ─────────────────────────
llm = Llama(model_path=MODEL_PATH)

# ───────────────────────── NEIS 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, month, api_key):
    from_ymd = f"{year}{month:02}01"
    to_ymd = f"{year}{month:02}31"
    url = f"https://open.neis.go.kr/hub/SchoolSchedule?KEY={api_key}&Type=json&pIndex=1&pSize=100&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()
    return data.get("SchoolSchedule", [{}])[1].get("row", [])

# ───────────────────────── GPT4ALL 응답 생성 ─────────────────────────
def generate_gpt_answer(user_question, events):
    context = "\n".join([f"{e['EVENT_NM']} : {e['AA_YMD']}" for e in events])
    prompt = f"""
### Instruction:
다음은 상리초등학교의 2024년 7월 학사일정입니다:
{context}
질문: "{user_question}"
위 학사일정 정보만 참고해서 질문에 해당하는 날짜를 찾아, 자연스럽고 간결한 한국어 문장으로 답해주세요. 예: "2024년 7월 24일입니다."
### Response:
"""

    result = ""
    for output in llm(prompt, stop=["###"], stream=True):
        result += output["choices"][0]["text"]
    return result.strip()

# ───────────────────────── 전체 처리 함수 ─────────────────────────
def answer_question(user_question):
    api_key = os.environ.get("NEIS_API_KEY", "a69e08342c8947b4a52cd72789a5ecaf")
    school_name = "상리초등학교"
    region_code = "R10"
    year = 2024
    month = 7

    school_code, region_code = get_school_info(region_code, school_name, api_key)
    events = get_schedule(region_code, school_code, year, month, api_key)

    if not events:
        return "학사일정 정보를 불러오지 못했습니다."

    return generate_gpt_answer(user_question, events)

# ───────────────────────── Gradio 인터페이스 ─────────────────────────
description = "📅 상리초등학교 2024년 7월 학사일정 기반 GPT 챗봇입니다. 예: '여름방학은 언제야?'"

gr.Interface(fn=answer_question,
             inputs=gr.Textbox(lines=2, placeholder="예: 여름방학식은 언제야?", label="질문 입력"),
             outputs=gr.Textbox(label="답변"),
             title="🎓 GPT4ALL 학사일정 Q&A",
             description=description).launch()