# app.py
import os, io, time, json, re, requests
from io import BytesIO   # 파일이 없어도, 메모리 속 데이터(바이트열) 를 파일처럼 다룰 수 있게 해줌. 웹에서 받은 이미지/파일을 저장하지 않고 바로 사용
from PIL import Image
import streamlit as st
from google import genai

st.set_page_config(page_title="멀티모달 LLM 데모 (Gemma 3 27B)", page_icon="🤖", layout="wide")

# ===== Sidebar: 설정 =====
st.sidebar.title("🔧 설정")
api_key = st.sidebar.text_input("GOOGLE_API_KEY", value=os.getenv("GOOGLE_API_KEY", ""), type="password")
model_name = st.sidebar.selectbox("모델", ["gemma-3-27b-it"], index=0)
rate_delay = st.sidebar.number_input("호출간 대기(초)", value=1.0, step=0.5, min_value=0.0)

if not api_key:
    st.warning("사이드바에 GOOGLE_API_KEY를 입력하거나 환경변수로 설정해 주세요.")
    st.stop()

client = genai.Client(api_key=api_key)

# ===== 유틸 =====
RETRY_LIMIT = 3

def load_image(source: str) -> Image.Image:
    """URL/로컬 모두 지원 + 투명 배경 보정."""
    session = requests.Session()
    session.headers.update({"User-Agent": "Mozilla/5.0"})
    last_err = None
    for _ in range(RETRY_LIMIT):
        try:
            if source.startswith(("http://", "https://")):
                resp = session.get(source, timeout=10)
                resp.raise_for_status()
                img = Image.open(BytesIO(resp.content))
            else:
                img = Image.open(source)

            if img.mode in ("RGBA", "LA", "P"):
                bg = Image.new("RGB", img.size, (255, 255, 255))
                img = img.convert("RGBA")
                bg.paste(img, mask=img.split()[-1] if len(img.split()) > 3 else None)
                img = bg
            else:
                img = img.convert("RGB")
            return img
        except Exception as e:
            last_err = e
            time.sleep(0.5)
    raise RuntimeError(f"이미지 로드 실패: {last_err}")

RESET_PROMPT = "이전 대화는 무시하고, 아래 지시에만 응답하세요.\n\n"

def try_parse_json(raw_text: str):
    """응답에서 JSON 추출 시도."""
    pattern = r'```(?:json)?\s*(\{[\s\S]*?\})\s*```|(\{[\s\S]*?\})'
    for g1, g2 in re.findall(pattern, raw_text):
        cand = (g1 or g2).strip()
        try:
            return json.loads(cand)
        except json.JSONDecodeError:
            pass
    cleaned = re.sub(r'```json|```', '', raw_text).strip()
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        return None

def infer_text(article: str, prompt: str):
    contents = [RESET_PROMPT + prompt.strip() + "\n\n" + article.strip()]
    resp = client.models.generate_content(model=model_name, contents=contents)
    time.sleep(rate_delay)
    text = (resp.text or "").strip()
    parsed = try_parse_json(text)
    return parsed if parsed is not None else {"text": text}

def infer_image(image: Image.Image, prompt: str):
    contents = [RESET_PROMPT + prompt.strip(), image]
    resp = client.models.generate_content(model=model_name, contents=contents)
    time.sleep(rate_delay)
    text = (resp.text or "").strip()
    parsed = try_parse_json(text)
    return parsed if parsed is not None else {"text": text}

# ===== 기본 프롬프트(간결 버전) =====
TEXT_PROMPT = (
    "당신은 기사 정보원 분석 전문가입니다. 다음 기사에서 정보원을 추출하고, "
    "정보원별 묘사 프레임을 긍정/중립/부정으로 판정해 JSON으로만 출력하세요.\n"
    '예시: {"sources": ["정보원A","정보원B"], "frames": ["중립","부정"]}'
)

IMAGE_PROMPT = (
    "당신은 보도사진 분석 전문가입니다. 제공된 사진에서 Donald Trump 존재 여부를 판단하고 "
    "감정(emotion: positive/negative/neutral)과 역동성(dynamism: high/medium/low)을 평가해 "
    'JSON으로만 출력하세요. Trump가 없으면 {"trump_present": false}만 반환하세요.'
)

# ===== UI =====
st.title("🤖 멀티모달 LLM 데모 (Gemma 3 27B)")
tab_text, tab_img = st.tabs(["📝 텍스트 분석", "🖼️ 이미지 분석"])

with tab_text:
    st.subheader("기사 텍스트 → 정보원 & 프레임 판정")
    article = st.text_area(
        "기사 본문", height=180,
        value="Donald Trump와 Nancy Pelosi가 회의장에서 격렬히 논쟁을 벌였다..."
    )
    user_prompt = st.text_area("프롬프트(옵션)", value=TEXT_PROMPT, height=120)
    if st.button("텍스트 분석 실행", type="primary", use_container_width=True):
        if not article.strip():
            st.error("기사 본문을 입력해 주세요.")
        else:
            with st.spinner("분석 중..."):
                try:
                    result = infer_text(article, user_prompt or TEXT_PROMPT)
                    st.success("완료")
                    st.json(result)
                    st.download_button("결과 JSON 다운로드", data=json.dumps(result, ensure_ascii=False, indent=2),
                                       file_name="text_result.json", mime="application/json")
                except Exception as e:
                    st.error(f"오류: {e}")

with tab_img:
    st.subheader("보도사진 → 인물 존재·감정·역동성 분석")
    col1, col2 = st.columns(2)
    with col1:
        img_url = st.text_input("이미지 URL")
    with col2:
        file = st.file_uploader("이미지 업로드", type=["jpg","jpeg","png","webp","gif","bmp"])

    img_prompt = st.text_area("프롬프트(옵션)", value=IMAGE_PROMPT, height=120)
    if st.button("이미지 분석 실행", type="primary", use_container_width=True):
        try:
            if file is not None:
                image = Image.open(io.BytesIO(file.read())).convert("RGB")
            elif img_url.strip():
                image = load_image(img_url.strip())
            else:
                st.error("이미지 URL을 입력하거나 파일을 업로드해 주세요.")
                st.stop()

            st.image(image, caption="입력 이미지", use_container_width=True)
            with st.spinner("분석 중..."):
                result = infer_image(image, img_prompt or IMAGE_PROMPT)
            st.success("완료")
            st.json(result)
            st.download_button("결과 JSON 다운로드", data=json.dumps(result, ensure_ascii=False, indent=2),
                               file_name="image_result.json", mime="application/json")
        except Exception as e:
            st.error(f"오류: {e}")