app.py

import json
import os
import re
import time
import urllib.parse
from typing import Any, Dict, List, Optional

import gradio as gr
import requests
import spaces
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import (
    AutoModelForVision2Seq,
    AutoProcessor,
)

# 최고 품질을 위한 대용량 모델 (ZeroGPU duration 최적화)
VL_MODEL_ID = "Qwen/Qwen2-VL-72B-Instruct"


def search_drug_web_simple(drug_name: str) -> str:
    """간단한 웹 검색으로 약물 정보 검증"""
    try:
        clean_name = re.sub(r'\(.*?\)|\d+mg|\d+mL|정|포|캡슐', '', drug_name).strip()
        sources = [
            f"https://www.health.kr/searchIdentity/search_result_detail.asp?searchStr={urllib.parse.quote(clean_name)}",
            f"https://terms.naver.com/search.naver?query={urllib.parse.quote(clean_name + ' 약')}"
        ]

        for url in sources:
            try:
                response = requests.get(url, timeout=3, headers={'User-Agent': 'Mozilla/5.0'})
                if response.status_code == 200 and len(response.text) > 1000:
                    text = response.text[:3000]
                    if any(kw in text for kw in ["효능", "효과", "복용", "주의"]):
                        return f"✓ 웹에서 {clean_name} 정보를 찾았습니다."
            except:
                continue
        return ""
    except:
        return ""


def _load_font():
    """한글 폰트 로드"""
    font_path = "NotoSansKR-Regular.ttf"
    if not os.path.exists(font_path):
        try:
            url = "https://github.com/notofonts/noto-cjk/raw/main/Sans/OTF/Korean/NotoSansKR-Regular.otf"
            response = requests.get(url)
            with open(font_path, "wb") as f:
                f.write(response.content)
        except:
            return None
    try:
        return ImageFont.truetype(font_path, 16)
    except:
        return None


DEFAULT_FONT = _load_font()


def _load_vl_model():
    """대용량 VL 모델 로드 - 최대 품질 + ZeroGPU 최적화"""
    device_map = "auto" if torch.cuda.is_available() else None

    # 8비트 양자화로 메모리 절약 (품질 유지하면서 메모리 1/2)
    model = AutoModelForVision2Seq.from_pretrained(
        VL_MODEL_ID,
        device_map=device_map,
        load_in_8bit=True,  # 8비트 양자화
        trust_remote_code=True,
    )

    processor = AutoProcessor.from_pretrained(VL_MODEL_ID, trust_remote_code=True)
    return model, processor


print("🔄 Loading Qwen2-VL-72B model with 8-bit quantization...")
VL_MODEL, VL_PROCESSOR = _load_vl_model()
print("✅ Model loaded successfully! (72B @ 8-bit)")


def _extract_assistant_content(decoded: str) -> str:
    """어시스턴트 응답 추출"""
    if "<|im_start|>assistant" in decoded:
        content = decoded.split("<|im_start|>assistant")[-1]
        content = content.replace("<|im_end|>", "").strip()
        return content
    return decoded.strip()


def _extract_json_block(text: str) -> Optional[str]:
    """JSON 블록 추출"""
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    return match.group(0)


def _sanitize_list(value: Any) -> List[str]:
    """리스트 정제"""
    if isinstance(value, (list, tuple)):
        return [str(v).strip() for v in value if str(v).strip()]
    if isinstance(value, str):
        return [v.strip() for v in re.split(r"[,;]", value) if v.strip()]
    return []


def _sanitize_medication(item: Dict[str, Any]) -> Dict[str, Any]:
    """약물 정보 정제"""
    def _to_str(val: Any) -> str:
        return "" if val is None else str(val).strip()

    times = item.get("times_per_day")
    if isinstance(times, (int, float)):
        times_str = str(int(times)) if float(times).is_integer() else str(times)
    else:
        times_str = _to_str(times)

    return {
        "name": _to_str(item.get("name")),
        "dose_per_intake": _to_str(item.get("dose_per_intake")),
        "times_per_day": times_str,
        "time_slots": _sanitize_list(item.get("time_slots")),
        "description": _to_str(item.get("description")),
        "efficacy": _to_str(item.get("efficacy")),
        "usage_precautions": _to_str(item.get("usage_precautions")),
        "side_effects": _to_str(item.get("side_effects")),
        "drug_interactions": _to_str(item.get("drug_interactions")),
        "warnings": _to_str(item.get("warnings")),
    }


def _parse_vl_response(text: str) -> Dict[str, Any]:
    """VL 모델 응답 파싱"""
    json_block = _extract_json_block(text)
    if not json_block:
        return {
            "raw_text": "",
            "medications": [],
            "warnings": ["모델 응답에서 JSON 형식을 찾지 못했습니다."],
        }

    try:
        data = json.loads(json_block)
    except json.JSONDecodeError:
        return {
            "raw_text": "",
            "medications": [],
            "warnings": ["JSON 파싱 실패"],
        }

    meds_raw = data.get("medications") or []
    medications = []
    if isinstance(meds_raw, list):
        for item in meds_raw:
            if isinstance(item, dict):
                medications.append(_sanitize_medication(item))

    warnings_raw = data.get("warnings")
    if isinstance(warnings_raw, list):
        warnings = [str(w).strip() for w in warnings_raw if str(w).strip()]
    elif warnings_raw:
        warnings = [str(warnings_raw).strip()]
    else:
        warnings = []

    return {
        "raw_text": str(data.get("raw_text", "")).strip(),
        "medications": medications,
        "warnings": warnings,
    }


@spaces.GPU(duration=120)  # 최대 2분 허용
def analyze_with_vl_model(image: Image.Image, task: str = "ocr") -> Any:
    """
    단일 VL 모델로 모든 작업 수행
    task: "ocr" (약봉투 분석) | "explain" (설명 생성) | "image_prompt" (이미지 프롬프트)
    """
    try:
        if task == "ocr":
            # 약봉투 OCR 및 정보 추출
            instructions = """사진 속 약봉투/처방전을 읽고 JSON 형식으로 답변하세요."""

            schema = """{
  "raw_text": "OCR로 읽은 전체 문장",
  "medications": [
    {
      "name": "약 이름 (상품명과 성분명)",
      "dose_per_intake": "1회 용량",
      "times_per_day": "하루 복용 횟수",
      "time_slots": ["복용 시간대"],
      "description": "약 설명",
      "efficacy": "이 약은 무엇입니까? (상세한 효능효과)",
      "usage_precautions": "이 약은 어떻게 복용합니까? (상세한 복용법)",
      "side_effects": "주요 부작용",
      "drug_interactions": "약물 상호작용",
      "warnings": "특별 주의사항"
    }
  ],
  "warnings": ["전체 경고"]
}"""

            messages = [
                {
                    "role": "system",
                    "content": "당신은 대한민국 약사입니다. 약봉투를 정확히 읽고 상세한 약물 정보를 제공합니다.",
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": instructions},
                        {"type": "text", "text": schema},
                        {"type": "image"},
                    ],
                },
            ]

            chat_text = VL_PROCESSOR.apply_chat_template(messages, add_generation_prompt=True)
            inputs = VL_PROCESSOR(text=[chat_text], images=[image], return_tensors="pt").to(VL_MODEL.device)

            output_ids = VL_MODEL.generate(
                **inputs,
                max_new_tokens=3072,
                temperature=0.3,
                top_p=0.95,
                do_sample=True,
            )

            decoded = VL_PROCESSOR.batch_decode(output_ids, skip_special_tokens=False)[0]
            assistant_text = _extract_assistant_content(decoded)
            return _parse_vl_response(assistant_text)

        elif task == "explain":
            # 설명 생성 (image는 None, text만 사용)
            return {"elderly_narrative": "", "child_narrative": "", "image_description": ""}

    except Exception as e:
        return {"error": str(e)}


def render_card(medications: List[Dict[str, Any]]) -> Image.Image:
    """현대적인 약물 카드 렌더링"""
    try:
        font_large = ImageFont.truetype("NotoSansKR-Regular.ttf", 28)
        font_medium = ImageFont.truetype("NotoSansKR-Regular.ttf", 20)
        font_small = ImageFont.truetype("NotoSansKR-Regular.ttf", 16)
    except:
        font_large = font_medium = font_small = None

    if not medications:
        canvas = Image.new("RGB", (900, 300), (255, 255, 255))
        draw = ImageDraw.Draw(canvas)
        draw.text((350, 130), "약 정보가 없습니다", fill=(140, 140, 140), font=font_medium)
        return canvas

    card_height_per_med = 240
    header_height = 120
    footer_height = 80
    total_height = header_height + (card_height_per_med * len(medications)) + footer_height

    width = 900
    canvas = Image.new("RGB", (width, total_height), (248, 250, 252))
    draw = ImageDraw.Draw(canvas)

    # 모던 헤더
    for i in range(header_height):
        alpha = i / header_height
        color = (
            int(99 + (248 - 99) * alpha),
            int(102 + (250 - 102) * alpha),
            int(241 + (252 - 241) * alpha),
        )
        draw.rectangle((0, i, width, i + 1), fill=color)

    draw.text((40, 35), "💊 복용 안내", fill=(30, 41, 59), font=font_large)
    draw.text((40, 75), f"{len(medications)}개 약품", fill=(71, 85, 105), font=font_small)

    y = header_height + 30

    for idx, med in enumerate(medications):
        card_y_start = y - 15
        card_y_end = y + 200

        # 카드 그림자
        draw.rounded_rectangle(
            (35, card_y_start + 5, width - 35, card_y_end + 5),
            radius=16,
            fill=(203, 213, 225),
        )

        # 카드 본체
        draw.rounded_rectangle(
            (30, card_y_start, width - 30, card_y_end),
            radius=16,
            fill=(255, 255, 255),
        )

        # 약 번호 배지
        badge_x, badge_y = 45, y
        draw.ellipse(
            (badge_x, badge_y, badge_x + 45, badge_y + 45),
            fill=(99, 102, 241),
        )
        draw.text((badge_x + 12, badge_y + 8), str(idx + 1), fill=(255, 255, 255), font=font_medium)

        # 약 이름
        name_text = med.get("name", "약 이름 미확인")
        draw.text((105, y + 8), name_text, fill=(15, 23, 42), font=font_medium)

        y += 60

        # 정보 섹션
        info_items = [
            ("📦", "용량", med.get('dose_per_intake', '-')),
            ("🔢", "횟수", f"{med.get('times_per_day', '-')}회/일"),
            ("🕐", "시간", ", ".join(med.get('time_slots') or ["-"])),
        ]

        for icon, label, value in info_items:
            draw.text((50, y), f"{icon} {label}", fill=(100, 116, 139), font=font_small)
            draw.text((160, y), value, fill=(30, 41, 59), font=font_small)
            y += 38

        y += 30

    # 푸터
    footer_y = total_height - footer_height + 25
    draw.text((40, footer_y), "※ 본 앱은 참고용이며, 실제 복약은 의사·약사의 지시를 따라주세요.",
              fill=(148, 163, 184), font=font_small)

    return canvas


def medications_to_csv(medications: List[Dict[str, Any]]) -> str:
    """CSV 생성"""
    if not medications:
        return ""

    rows = ["약명,1회용량,1일횟수,시간대"]
    for med in medications:
        row = [
            med.get("name", ""),
            med.get("dose_per_intake", ""),
            med.get("times_per_day", ""),
            ";".join(med.get("time_slots") or []),
        ]
        rows.append(",".join(row))

    return "\n".join(rows)


def format_warnings(warnings: List[str]) -> str:
    """경고 메시지 포맷"""
    if not warnings:
        return "✅ 인식된 정보가 충분합니다."

    lines = ["### ⚠️ 확인 필요"]
    for warn in warnings:
        lines.append(f"- {warn}")
    lines.append("\n> 의료진의 지시가 가장 정확합니다.")
    return "\n".join(lines)


@spaces.GPU(duration=90)  # 설명 생성은 90초
def generate_full_explanation(medications: List[Dict[str, Any]], raw_text: str, web_info: str = "") -> Dict[str, str]:
    """VL 모델로 설명 생성"""
    try:
        med_summary = "\n".join([
            f"- {med.get('name')} {med.get('dose_per_intake')} (하루 {med.get('times_per_day')}회)"
            for med in medications
        ])

        web_context = f"\n\n웹 검증: {web_info}" if web_info else ""

        prompt = f"""다음 약물 정보를 바탕으로 어르신과 어린이를 위한 설명을 작성하세요.

약 정보:
{med_summary}

원문: {raw_text}{web_context}

JSON 형식으로 답변:
{{
  "elderly": {{
    "narrative": "어르신을 위한 설명 (존댓말, 구체적, 5-7문장)",
    "image_description": "약 복용 장면 묘사 (한국어)"
  }},
  "child": {{
    "narrative": "어린이를 위한 설명 (쉬운 말, 재미있게, 4-6문장)",
    "image_description": "약 복용 장면 묘사 (한국어)"
  }}
}}"""

        messages = [
            {
                "role": "system",
                "content": "당신은 20년 경력 임상약사입니다. 환자 교육 전문가입니다.",
            },
            {
                "role": "user",
                "content": prompt,
            },
        ]

        chat_text = VL_PROCESSOR.apply_chat_template(messages, add_generation_prompt=True)
        inputs = VL_PROCESSOR(text=[chat_text], images=None, return_tensors="pt").to(VL_MODEL.device)

        output_ids = VL_MODEL.generate(
            **inputs,
            max_new_tokens=2048,
            temperature=0.8,
            top_p=0.92,
            do_sample=True,
        )

        decoded = VL_PROCESSOR.batch_decode(output_ids, skip_special_tokens=False)[0]
        text = _extract_assistant_content(decoded)

        json_block = _extract_json_block(text)
        if json_block:
            data = json.loads(json_block)
            elderly = data.get("elderly", {})
            child = data.get("child", {})

            return {
                "elderly_narrative": str(elderly.get("narrative", "")).strip(),
                "child_narrative": str(child.get("narrative", "")).strip(),
            }

        return {
            "elderly_narrative": "설명을 생성하지 못했습니다.",
            "child_narrative": "설명을 생성하지 못했습니다.",
        }

    except Exception as e:
        return {
            "elderly_narrative": "설명 생성 중 오류 발생",
            "child_narrative": "설명 생성 중 오류 발생",
        }


def run_pipeline(image: Optional[Image.Image], progress=gr.Progress()):
    """메인 파이프라인"""
    if image is None:
        return (
            "이미지를 업로드하세요.",
            None,
            None,
            "이미지를 먼저 업로드해 주세요.",
            "📷 약 봉투 사진을 올리면 인식이 시작됩니다.",
            "",
            "약물 정보가 표시됩니다.",
        )

    progress(0, desc="🔍 약봉투 이미지 분석 중...")
    result = analyze_with_vl_model(image, task="ocr")

    medications = result.get("medications") or []

    # 웹 검색
    progress(0.25, desc="🌐 웹에서 약물 정보 검증 중...")
    web_info_results = []
    for med in medications[:3]:
        drug_name = med.get("name", "")
        if drug_name:
            web_info = search_drug_web_simple(drug_name)
            if web_info:
                web_info_results.append(web_info)
                med["web_verified"] = True

    web_search_info = "\n".join(web_info_results) if web_info_results else ""

    progress(0.5, desc="💬 설명 생성 중...")
    narratives = generate_full_explanation(medications, result.get("raw_text", ""), web_search_info)

    progress(0.75, desc="🎨 카드 렌더링 중...")
    card_img = render_card(medications)
    csv_row = medications_to_csv(medications)

    # 상세 정보
    detailed_info = "# 💊 약물 상세 정보\n\n"

    if web_search_info:
        detailed_info += "✅ **웹 검증 완료**\n\n"
        detailed_info += f"> {web_search_info}\n\n---\n\n"

    for idx, med in enumerate(medications):
        web_badge = " 🌐" if med.get("web_verified") else ""
        detailed_info += f"## {idx + 1}. {med.get('name', '약 이름 미확인')}{web_badge}\n\n"

        if med.get("efficacy"):
            detailed_info += f"### 🔍 이 약은 무엇입니까?\n{med.get('efficacy')}\n\n"

        if med.get("usage_precautions"):
            detailed_info += f"### 📋 이 약은 어떻게 복용합니까?\n{med.get('usage_precautions')}\n\n"

        if med.get("side_effects"):
            detailed_info += f"### ⚠️ 부작용\n{med.get('side_effects')}\n\n"

        if med.get("drug_interactions"):
            detailed_info += f"### 🔄 약물 상호작용\n{med.get('drug_interactions')}\n\n"

        if med.get("warnings"):
            detailed_info += f"### ⚡ 특별 주의사항\n{med.get('warnings')}\n\n"

        detailed_info += "---\n\n"

    # 설명 마크다운
    markdown = (
        "## 👴 어르신을 위한 설명\n\n"
        + (narratives.get("elderly_narrative") or "설명을 준비하지 못했습니다.")
        + "\n\n## 👶 어린이를 위한 설명\n\n"
        + (narratives.get("child_narrative") or "설명을 준비하지 못했습니다.")
        + "\n\n> 💡 항상 의료진의 안내를 우선하세요."
    )

    warnings_md = format_warnings(result.get("warnings", []))
    raw_text = result.get("raw_text", "")
    json_text = json.dumps(result, ensure_ascii=False, indent=2)

    progress(1.0, desc="✅ 완료!")
    return json_text, card_img, csv_row, markdown, warnings_md, raw_text, detailed_info


# 현대적인 CSS
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');

:root {
    --primary: #6366f1;
    --primary-dark: #4f46e5;
    --secondary: #8b5cf6;
    --success: #10b981;
    --warning: #f59e0b;
    --danger: #ef4444;
    --gray-50: #f9fafb;
    --gray-100: #f3f4f6;
    --gray-200: #e5e7eb;
    --gray-300: #d1d5db;
    --gray-600: #4b5563;
    --gray-800: #1f2937;
    --gray-900: #111827;
}

body {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
}

.gradio-container {
    max-width: 1400px !important;
    margin: auto;
    background: rgba(255, 255, 255, 0.95);
    border-radius: 24px;
    box-shadow: 0 25px 50px -12px rgba(0, 0, 0, 0.25);
    padding: 40px;
}

.hero-section {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 20px;
    padding: 50px 40px;
    margin-bottom: 40px;
    color: white;
    box-shadow: 0 20px 40px -10px rgba(102, 126, 234, 0.4);
}

.hero-section h1 {
    font-size: 2.5rem;
    font-weight: 700;
    margin-bottom: 16px;
    text-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
}

.hero-section p {
    font-size: 1.15rem;
    opacity: 0.95;
    line-height: 1.6;
}

.card {
    background: white;
    border-radius: 16px;
    padding: 32px;
    box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
    transition: all 0.3s ease;
}

.card:hover {
    box-shadow: 0 20px 25px -5px rgba(0, 0, 0, 0.1), 0 10px 10px -5px rgba(0, 0, 0, 0.04);
    transform: translateY(-2px);
}

.primary-btn button {
    background: linear-gradient(135deg, var(--primary) 0%, var(--secondary) 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 600 !important;
    font-size: 1.05rem !important;
    padding: 16px 32px !important;
    border-radius: 12px !important;
    box-shadow: 0 10px 20px -5px rgba(99, 102, 241, 0.4) !important;
    transition: all 0.3s ease !important;
}

.primary-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 15px 30px -5px rgba(99, 102, 241, 0.5) !important;
}

.tab-nav button {
    font-weight: 500 !important;
    border-radius: 8px !important;
    transition: all 0.2s ease !important;
}

.tab-nav button.selected {
    background: var(--primary) !important;
    color: white !important;
}

.notice {
    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%);
    border-left: 4px solid var(--warning);
    border-radius: 12px;
    padding: 20px;
    color: var(--gray-800);
}

.output-card {
    background: var(--gray-50);
    border-radius: 16px;
    padding: 28px;
    border: 1px solid var(--gray-200);
}

.gr-image {
    border-radius: 16px !important;
    box-shadow: 0 10px 20px -5px rgba(0, 0, 0, 0.1) !important;
}

.csv-box textarea {
    font-family: 'JetBrains Mono', 'Courier New', monospace !important;
    font-size: 0.9rem !important;
    background: var(--gray-900) !important;
    color: #10b981 !important;
    border-radius: 12px !important;
}

.accordion {
    border-radius: 12px !important;
    border: 1px solid var(--gray-200) !important;
}

h1, h2, h3 {
    font-weight: 600;
    color: var(--gray-900);
}

.markdown-text {
    line-height: 1.8;
    color: var(--gray-800);
}
"""

HERO_HTML = """
<div class="hero-section">
    <h1>🏥 MedCard Pro</h1>
    <p>
        <strong>AI 기반 스마트 약물 관리 시스템</strong><br>
        Qwen2-VL-72B (8비트 최적화)가 약봉투를 최고 정확도로 분석하고,<br>
        웹에서 실시간으로 정보를 검증하여 프로페셔널한 복약 안내를 제공합니다.
    </p>
</div>
"""

# Gradio 인터페이스
with gr.Blocks(theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
    gr.HTML(HERO_HTML)

    with gr.Row():
        with gr.Column(scale=5, elem_classes=["card"]):
            gr.Markdown("### 📸 약 봉투 사진 업로드")
            img_in = gr.Image(type="pil", label="약봉투/처방전 사진", height=400)
            warn_md = gr.Markdown("💡 약 봉투 사진을 올려주세요. AI가 자동으로 분석합니다.", elem_classes=["notice"])
            btn = gr.Button("🚀 분석 시작", elem_classes=["primary-btn"], size="lg")

        with gr.Column(scale=7, elem_classes=["card"]):
            gr.Markdown("### 📊 분석 결과")

            with gr.Tabs():
                with gr.Tab("📚 약물 상세 정보"):
                    detailed_info_md = gr.Markdown("분석을 시작하면 여기에 약물 정보가 표시됩니다.", elem_classes=["output-card"])

                with gr.Tab("👥 쉬운 설명"):
                    explain_md = gr.Markdown("어르신과 어린이를 위한 설명이 표시됩니다.", elem_classes=["output-card"])

                with gr.Tab("📅 복용 일정"):
                    card_out = gr.Image(type="pil", label="일정 카드")

            with gr.Accordion("🔍 상세 분석 결과", open=False):
                raw_box = gr.Textbox(label="OCR 원문", lines=4, interactive=False)
                csv_box = gr.Textbox(label="CSV 데이터", lines=3, elem_classes=["csv-box"])
                json_out = gr.Code(label="JSON 데이터", language="json")

    btn.click(
        run_pipeline,
        inputs=img_in,
        outputs=[json_out, card_out, csv_box, explain_md, warn_md, raw_box, detailed_info_md],
    )

    gr.Markdown(
        """
        ---

        ### ℹ️ 주의사항

        이 서비스는 **참고용 도구**입니다. 실제 복약은 반드시 **의사·약사의 지시**에 따라주세요.

        🔒 개인정보는 저장되지 않으며, 모든 처리는 실시간으로 이루어집니다.
        """
    )

if __name__ == "__main__":
    demo.queue().launch()