import subprocess

# 필요한 라이브러리 설치 및 업데이트
subprocess.run(["pip", "install", "--upgrade", "pip"])
subprocess.run(["pip", "install", "--upgrade", "openai", "yfinance", "gradio", "matplotlib", "Pillow"])

# 이후 코드 실행
import requests
import gradio as gr
import yfinance as yf
import matplotlib.pyplot as plt
import io
from PIL import Image
from datetime import datetime, timedelta
from openai import OpenAI
import os

# Perplexity AI API 설정
API_KEY = "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"  # 여기에 Perplexity AI API 키를 입력하세요.

def get_real_news_summary(company, date):
    # OpenAI 클라이언트 초기화
    client = OpenAI(api_key=API_KEY, base_url="https://api.perplexity.ai")
    
    # API 요청을 위한 메시지 구성
    messages = [
        {"role": "system", "content": "You are a helpful assistant that summarizes stock news in Korean."},
        {"role": "user", "content": f"Summarize the latest stock news for {company} on {date} in Korean. If there's no specific news for that date, provide the most recent relevant information."}
    ]
    
    try:
        # API 요청
        response = client.chat.completions.create(
            model="llama-3.1-sonar-large-128k-online",
            messages=messages
        )
        
        # 응답에서 요약 추출
        summary = response.choices[0].message.content
        return summary
    except Exception as e:
        return f"뉴스 요약 중 에러가 발생했습니다: {str(e)}"

# 뉴스 요약을 가져오는 함수
def handle_click(input_value, date_clicked):
    return get_real_news_summary(input_value, date_clicked)

# Gradio에서 사용할 함수 (뉴스 요약 포함)
def update_news(input_value, selected_date):
    if selected_date == "" or selected_date is None:
        return "날짜를 선택해주세요."
    else:
        return handle_click(input_value, selected_date)

# 종목명과 티커를 매핑하는 딕셔너리
name_to_ticker = {
    "삼성전자": "005930.KS",
    "SK바이오팜": "326030.KS",
    "Apple": "AAPL",
    # 필요한 종목들을 추가하세요
}

# 주가 데이터를 가져오고 조건에 맞는 날짜와 그래프를 반환하는 함수
def display_stock_with_highlight(input_value, change_type, percent_change):
    try:
        ticker = name_to_ticker.get(input_value, input_value)
        stock = yf.Ticker(ticker)
        stock_data = stock.history(period="5y")  # 최근 5년 데이터로 제한

        if stock_data.empty:
            return "주가 데이터를 찾을 수 없습니다.", []

        stock_data['Change'] = stock_data['Close'].pct_change() * 100

        percent_change = float(percent_change)

        if change_type == "상승":
            highlight_data = stock_data[stock_data['Change'] >= percent_change]
            color = "darkorange"
        elif change_type == "하락":
            highlight_data = stock_data[stock_data['Change'] <= -percent_change]
            color = "purple"
        else:
            return "Invalid change type", []

        dates = stock_data.index.to_numpy()
        closing_prices = stock_data['Close'].to_numpy()

        plt.figure(figsize=(10, 6))
        plt.plot(dates, closing_prices, color='gray', label=input_value)
        plt.scatter(highlight_data.index, highlight_data['Close'], color=color, label=f'{change_type} 포인트')

        for index, row in highlight_data.iterrows():
            plt.text(index, row['Close'], index.strftime('%Y-%m-%d'), fontsize=10, fontweight='bold', color=color, ha='right')
            plt.axvline(x=index, color=color, linestyle='--', linewidth=1)  # x축과의 연결선 점선으로 표시

        plt.title(f'{input_value} 주가 추이 ({change_type} 표시)')
        plt.xlabel('날짜')
        plt.ylabel('종가')
        plt.legend()

        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        plt.close()
        buf.seek(0)
        img = Image.open(buf)

        highlight_dates = highlight_data.index.strftime('%Y-%m-%d').tolist()

        return img, gr.update(choices=highlight_dates)
    except Exception as e:
        return f"Error processing data: {e}", gr.update(choices=[])

# Gradio 인터페이스 생성 (3열 레이아웃)
with gr.Blocks() as demo:
    gr.Markdown("## 주가 그래프와 뉴스 요약")

    with gr.Row():
        with gr.Column():  # 입력값을 담을 첫 번째 열
            input_value = gr.Textbox(label="종목명 또는 티커 입력", placeholder="예: SK바이오팜, AAPL")
            change_type = gr.Dropdown(choices=["상승", "하락"], label="상승 또는 하락 선택", value="상승")
            percent_change = gr.Textbox(label="변동 퍼센트 (%)", placeholder="예: 10", value="10")

            submit_btn = gr.Button("Submit")

            # 예제 (이전 레이아웃으로 복원)
            examples = [["SK바이오팜"], 
                        ["나스닥 시총 1위"], 
                        ["나스닥 제약주 시총 1위"], 
                        ["나스닥 바이오텍 시총 1위"], 
                        ["코스피 시총 1위"], 
                        ["코스닥 시총 1위"]]
            gr.Examples(examples=examples, inputs=[input_value])

        with gr.Column():  # 그래프를 출력할 두 번째 열
            plot = gr.Image(label="주가 그래프")
            date_dropdown = gr.Dropdown(label="조건에 해당하는 날짜 선택", choices=[])

        with gr.Column():  # 뉴스 요약을 출력할 세 번째 열
            news_output = gr.Textbox(label="뉴스 요약")

    # Submit 버튼 클릭 시 그래프 및 날짜 드롭다운 업데이트
    submit_btn.click(
        fn=display_stock_with_highlight,
        inputs=[input_value, change_type, percent_change],
        outputs=[plot, date_dropdown]
    )

    # 날짜 선택 시 뉴스 요약 업데이트
    date_dropdown.change(
        fn=update_news,
        inputs=[input_value, date_dropdown],
        outputs=[news_output]
    )

# Gradio 실행
demo.launch()