extractor.py

"""
Upstage Info Extract 연동 모듈
가격, 단위, 날짜 등 정보 추출 (Solar LLM 기반)
"""

import os
import re
from dataclasses import dataclass

from dotenv import load_dotenv

load_dotenv()


@dataclass
class ExtractedPrice:
    """추출된 가격 정보"""
    original: str  # 원본 표기 (예: "$19.99")
    value: float  # 숫자 값
    currency: str  # 통화 (USD, EUR, KRW 등)
    korean_format: str  # 한국식 표기 (예: "₩26,000")


@dataclass
class ExtractedDate:
    """추출된 날짜 정보"""
    original: str  # 원본 표기 (예: "12/25/2024")
    korean_format: str  # 한국식 표기 (예: "2024년 12월 25일")


@dataclass
class ExtractedUnit:
    """추출된 단위 정보"""
    original: str  # 원본 표기 (예: "5.2 miles")
    value: float  # 숫자 값
    unit: str  # 원본 단위
    korean_value: float  # 변환된 값
    korean_unit: str  # 한국식 단위
    korean_format: str  # 한국식 표기 (예: "8.4km")


class InfoExtractor:
    """텍스트에서 현지화 필요 정보 추출"""

    # 통화 변환 (대략적인 환율)
    CURRENCY_TO_KRW = {
        "USD": 1300,
        "EUR": 1400,
        "GBP": 1650,
        "JPY": 9,
        "CNY": 180,
    }

    # 단위 변환
    UNIT_CONVERSIONS = {
        "mile": ("km", 1.60934),
        "miles": ("km", 1.60934),
        "ft": ("m", 0.3048),
        "feet": ("m", 0.3048),
        "inch": ("cm", 2.54),
        "inches": ("cm", 2.54),
        "lb": ("kg", 0.453592),
        "lbs": ("kg", 0.453592),
        "oz": ("g", 28.3495),
        "gallon": ("L", 3.78541),
        "gallons": ("L", 3.78541),
        "°F": ("°C", None),  # 특수 변환
    }

    def __init__(self):
        self.api_key = os.getenv("UPSTAGE_API_KEY")
        self.base_url = "https://api.upstage.ai/v1/solar/chat/completions"

    def extract_prices(self, text: str) -> list[ExtractedPrice]:
        """텍스트에서 가격 정보를 추출하고 한국 원화로 변환합니다.

        Args:
            text: 분석할 텍스트

        Returns:
            추출된 가격 정보 리스트
        """
        prices = []

        # 가격 패턴들
        patterns = [
            (r'\$(\d+(?:,\d{3})*(?:\.\d{2})?)', 'USD'),
            (r'€(\d+(?:,\d{3})*(?:\.\d{2})?)', 'EUR'),
            (r'£(\d+(?:,\d{3})*(?:\.\d{2})?)', 'GBP'),
            (r'¥(\d+(?:,\d{3})*)', 'JPY'),
            (r'(\d+(?:,\d{3})*(?:\.\d{2})?)\s*(?:USD|dollars?)', 'USD'),
            (r'(\d+(?:,\d{3})*(?:\.\d{2})?)\s*(?:EUR|euros?)', 'EUR'),
        ]

        for pattern, currency in patterns:
            matches = re.finditer(pattern, text, re.IGNORECASE)
            for match in matches:
                value_str = match.group(1).replace(',', '')
                value = float(value_str)

                # 원화 변환
                krw_value = int(value * self.CURRENCY_TO_KRW.get(currency, 1))
                korean_format = f"₩{krw_value:,}"

                prices.append(ExtractedPrice(
                    original=match.group(0),
                    value=value,
                    currency=currency,
                    korean_format=korean_format
                ))

        return prices

    def extract_dates(self, text: str) -> list[ExtractedDate]:
        """텍스트에서 날짜 정보를 추출하고 한국식으로 변환합니다.

        Args:
            text: 분석할 텍스트

        Returns:
            추출된 날짜 정보 리스트
        """
        dates = []

        # 날짜 패턴들 (MM/DD/YYYY, Month DD, YYYY 등)
        patterns = [
            # MM/DD/YYYY or MM-DD-YYYY
            (r'(\d{1,2})[/\-](\d{1,2})[/\-](\d{4})', 'MDY'),
            # YYYY/MM/DD or YYYY-MM-DD
            (r'(\d{4})[/\-](\d{1,2})[/\-](\d{1,2})', 'YMD'),
            # Month DD, YYYY
            (r'(Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:tember)?|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)\s+(\d{1,2}),?\s+(\d{4})', 'MONTH'),
        ]

        month_map = {
            'jan': 1, 'january': 1, 'feb': 2, 'february': 2,
            'mar': 3, 'march': 3, 'apr': 4, 'april': 4,
            'may': 5, 'jun': 6, 'june': 6, 'jul': 7, 'july': 7,
            'aug': 8, 'august': 8, 'sep': 9, 'september': 9,
            'oct': 10, 'october': 10, 'nov': 11, 'november': 11,
            'dec': 12, 'december': 12
        }

        for pattern, fmt in patterns:
            matches = re.finditer(pattern, text, re.IGNORECASE)
            for match in matches:
                if fmt == 'MDY':
                    month, day, year = int(match.group(1)), int(match.group(2)), int(match.group(3))
                elif fmt == 'YMD':
                    year, month, day = int(match.group(1)), int(match.group(2)), int(match.group(3))
                elif fmt == 'MONTH':
                    month = month_map.get(match.group(1).lower(), 1)
                    day, year = int(match.group(2)), int(match.group(3))

                korean_format = f"{year}년 {month}월 {day}일"

                dates.append(ExtractedDate(
                    original=match.group(0),
                    korean_format=korean_format
                ))

        return dates

    def extract_units(self, text: str) -> list[ExtractedUnit]:
        """텍스트에서 단위 정보를 추출하고 한국식으로 변환합니다.

        Args:
            text: 분석할 텍스트

        Returns:
            추출된 단위 정보 리스트
        """
        units = []

        for unit, (korean_unit, factor) in self.UNIT_CONVERSIONS.items():
            pattern = rf'(\d+(?:\.\d+)?)\s*{re.escape(unit)}'
            matches = re.finditer(pattern, text, re.IGNORECASE)

            for match in matches:
                value = float(match.group(1))

                # 화씨 -> 섭씨 특수 변환
                if unit == "°F":
                    korean_value = round((value - 32) * 5 / 9, 1)
                else:
                    korean_value = round(value * factor, 1)

                korean_format = f"{korean_value}{korean_unit}"

                units.append(ExtractedUnit(
                    original=match.group(0),
                    value=value,
                    unit=unit,
                    korean_value=korean_value,
                    korean_unit=korean_unit,
                    korean_format=korean_format
                ))

        return units

    def extract_all(self, text: str) -> dict:
        """텍스트에서 모든 현지화 필요 정보를 추출합니다.

        Args:
            text: 분석할 텍스트

        Returns:
            추출된 모든 정보 딕셔너리
        """
        return {
            "prices": [vars(p) for p in self.extract_prices(text)],
            "dates": [vars(d) for d in self.extract_dates(text)],
            "units": [vars(u) for u in self.extract_units(text)],
        }