theme_info.py

import argparse
import csv
import html
import json
import os

from openai import OpenAI
client = OpenAI()


def parse_stock_input(stock_arg):
    """입력 문자열 → 종목 리스트 변환"""
    if not isinstance(stock_arg, str):
        return []
    items = stock_arg.split(",")
    result = []
    for item in items:
        s = item.strip()
        if (s.startswith('"') and s.endswith('"')) or (s.startswith("'") and s.endswith("'")):
            s = s[1:-1]
        if s:
            result.append(s)
    return result


class theme_info:
    def __init__(self, data):
        self.path = data
        self._store = {}
        self._raw = []

        if not os.path.exists(data):
            raise FileNotFoundError(f"Data file not found: {data}")

        # CSV 로드
        with open(data, "r", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            required_cols = {"NAME", "THEME", "THEME_2", "THEME_3"}
            headers = set(n.strip() for n in (reader.fieldnames or []))

            if not required_cols.issubset(headers):
                raise ValueError(
                    f"CSV must contain headers {sorted(required_cols)}. Found: {reader.fieldnames}"
                )

            for row in reader:
                name = (row.get("NAME") or "").strip()
                theme = html.unescape((row.get("THEME") or "").strip())
                theme2 = html.unescape((row.get("THEME_2") or "").strip())
                theme3 = html.unescape((row.get("THEME_3") or "").strip())

                desc = f"Theme: {theme} | Theme 2: {theme2} | Theme 3: {theme3}"
                self._store[name] = desc

                self._raw.append({
                    "NAME": name,
                    "THEME": theme,
                    "THEME_2": theme2,
                    "THEME_3": theme3
                })

        self._ci_index = {name.lower(): name for name in self._store.keys()}

    def get_ticker_and_name(self, stock):
        """GPT로 티커와 공식 종목명 조회"""
        prompt = f"""
                아래 종목명에 대해 Yahoo Finance 기준:
                1) 티커(symbol)
                2) 공식 종목명(full name)

                단, 중요한 조건:
                - 사용자가 입력한 종목명이 한국어라면 → 공식 종목명도 한국어로 반환하세요.
                - 사용자가 입력한 종목명에 한국어가 조금이라도 포함되어 있다면 → 공식 종목명도 한국어로 반환하세요. 
                - 사용자가 입력한 종목명이 영어라면 → 공식 종목명도 영어로 반환하세요.
                - 한국 종목은 '.KS', '.KQ' 등 KRX 형식의 티커.
                - 해외 종목(미국/유럽/홍콩 등)은 해당 시장의 일반 티커를 사용.

                예시:
                사용자 입력: 'SK하이닉스' → {{
                  "ticker": "000660.KS",
                  "name": "SK하이닉스"
                }}
                사용자 입력: 'SKHynix' → {{
                  "ticker": "000660.KS",
                  "name": "SK Hynix Inc."
                }}
                사용자 입력: '현대자동차' → {{
                  "ticker": "000660.KS",
                  "name": "현대차"
                }}

                종목명: {stock}

                오직 JSON으로만 답변하세요.
                """
        resp = client.responses.create(model="gpt-5.1", input=prompt)
        try:
            parsed = json.loads(resp.output_text)
            return parsed.get("ticker"), parsed.get("name")
        except:
            print("[ERROR] GPT 티커/풀네임 JSON 파싱 실패")
            return None, None

    def find_existing_row(self, full_name):
        """CSV 내 기존 종목 여부 확인"""
        return next(
            (r for r in self._raw if r["NAME"].lower() == full_name.lower()),
            None
        )

    def generate_theme_gpt(self, stock_name):
        """GPT로 테마 3개 생성"""
        prompt = f"""
                종목 '{stock_name}'과 관련된 투자 테마 3개 추천
                괄호 없이 간단하게 핵심 테마 표현
                - 사용자가 입력한 종목명이 한국어라면 → 테마도 한국어로 반환하세요.
                - 사용자가 입력한 종목명에 한국어가 조금이라도 포함되어 있다면 → 테마도 한국어로 반환하세요. 
                - 사용자가 입력한 종목명이 영어라면 → 테마도 영어로 반환하세요.
                JSON 형식으로 출력: {{"theme": "...", "theme2": "...", "theme3": "..."}}
                """

        try:
            response = client.responses.create(model="gpt-5.1", input=prompt)
            return json.loads(response.output_text)
        except:
            print(f"[WARN] GPT 테마 생성 실패({stock_name})")
            return {"theme": "", "theme2": "", "theme3": ""}

    def _save_to_csv(self, row):
        """단일 row를 CSV에 append 저장"""
        file_exists = os.path.exists(self.path)
        with open(self.path, "a", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(
                f,
                fieldnames=["NAME", "THEME", "THEME_2", "THEME_3"],
                quoting=csv.QUOTE_ALL
            )
            if not file_exists:
                writer.writeheader()
            writer.writerow(row)
        print(f"CSV append 완료: {row['NAME']}")

    def make(self, stock):
        """종목명 → 티커/공식명 조회 → 기존 CSV 확인 → GPT로 테마 생성 → CSV 저장"""

        # 1) 티커 + 공식명 조회
        ticker_symbol, full_name = self.get_ticker_and_name(stock)
        if not ticker_symbol:
            print("[ERROR] 티커 조회 실패")
            return None

        print(f"GPT 결과 → ticker: {ticker_symbol}, full name: {full_name}")

        # 2) 기존 CSV 체크
        existing = self.find_existing_row(full_name)
        if existing:
            print(f"'{full_name}' 기존 CSV 존재 → 기존값 반환")
            return [existing]

        # 3) GPT로 테마 생성
        theme_data = self.generate_theme_gpt(full_name)
        theme = theme_data.get("theme", "").strip()
        theme2 = theme_data.get("theme2", "").strip()
        theme3 = theme_data.get("theme3", "").strip()

        new_row = {
            "NAME": full_name,
            "THEME": theme,
            "THEME_2": theme2,
            "THEME_3": theme3
        }

        # CSV 저장
        self._save_to_csv(new_row)

        # _store 업데이트
        desc = f"Theme: {new_row['THEME']} | Theme 2: {new_row['THEME_2']} | Theme 3: {new_row['THEME_3']}"
        self._store[full_name] = desc
        self._raw.append(new_row)

        print("[STEP] make() 완료")
        return [new_row]

    def get(self, stock):
        """종목명으로 description 조회"""
        if not stock or not isinstance(stock, str):
            return None
        key = stock.strip()
        if key in self._store:
            return self._store[key]
        canonical = self._ci_index.get(key.lower())
        if canonical:
            return self._store[canonical]
        # 부분 일치 1개만 반환
        candidates = [n for n in self._store if key.lower() in n.lower()]
        if len(candidates) == 1:
            return self._store[candidates[0]]
        return None


def ensure_parent_dir(json_path):
    """파일 저장 전 상위 디렉토리 생성"""
    parent = os.path.dirname(os.path.abspath(json_path))
    if parent and not os.path.exists(parent):
        os.makedirs(parent, exist_ok=True)


def append_to_json(records, json_path, class_name="theme_info"):
    """records → JSON append 저장"""
    ensure_parent_dir(json_path)
    if os.path.exists(json_path):
        try:
            with open(json_path, "r", encoding="utf-8") as rf:
                data = json.load(rf)
        except:
            print("[WARN] 기존 JSON 로드 실패. 새 파일 생성")
            data = {class_name: []}
    else:
        data = {class_name: []}

    data.setdefault(class_name, [])
    data[class_name].extend(records)

    with open(json_path, "w", encoding="utf-8") as wf:
        json.dump(data, wf, ensure_ascii=False, indent=2)
    print(f"[5/5] {json_path} 업데이트 완료.\n")


def main():
    parser = argparse.ArgumentParser(description="종목의 테마 검색 도구")
    parser.add_argument("--stock", default="Meta", nargs="+", help="종목 입력. 여러 개는 공백 또는 콤마(,) 로 구분 가능")
    parser.add_argument("--data", default="/work/portfolio/data/theme_info.csv")
    parser.add_argument("--output", default="output.json")
    args = parser.parse_args()

    # --- 1. args 확인 ---
    stock_str = " ".join(args.stock)
    stock_list = parse_stock_input(stock_str)
    if not stock_list:
        print("[ERROR] --stock에 유효한 종목 없음")
        return
    print("[1/5] 인자(args) 확인.")
    print(f"[INFO] --stock {args.stock}")
    print(f"[INFO] --data {args.data}")
    print(f"[INFO] --output {args.output}")
    print(f"[INFO] 입력된 종목 리스트: {stock_list}\n")

    # --- 2. 객체 생성 ---
    ti = theme_info(args.data)
    print("[2/5] theme_info 객체 생성 완료.\n")

    # --- 3. make 함수 실행 ---
    print("[3/5] make() 실행.")
    make_results = []
    for stock in stock_list:
        print(f"[INFO] 처리 중: {stock}")
        rows = ti.make(stock)
        if rows:
            make_results.extend(rows)
    print()

    # --- 4. get 함수 실행 ---
    print("[4/5] get() 실행.")
    all_records = []
    for row in make_results:
        official_name = row["NAME"]
        desc = ti.get(official_name)
        all_records.append({"stock": official_name, "desc": desc})
        print(f"[INFO] {official_name} 조회 완료")
        print(f"{desc}")
    print()

    # --- 5. 저장 ---
    if all_records:
        append_to_json(all_records, args.output, class_name="theme_info")
    else:
        print("[INFO] 저장할 JSON 데이터 없음")


if __name__ == "__main__":
    main()