import gradio as gr
import json
import os
from pathlib import Path
import requests
import base64
from io import BytesIO

# Load watchlist configuration
def load_watchlists():
    """Load watchlist configuration from JSON file"""
    config_path = Path("watchlists_config.json")
    if config_path.exists():
        with open(config_path, 'r') as f:
            return json.load(f)
    return {}

def load_stocks_from_file(filename):
    """Load stock symbols from a watchlist file"""
    file_path = Path(filename)
    if file_path.exists():
        with open(file_path, 'r') as f:
            content = f.read().strip()
            # Split by comma and clean whitespace
            stocks = [s.strip() for s in content.split(',') if s.strip()]
            return stocks
    return []

def generate_chart_url(symbol, timeframe):
    """Generate StockCharts URL based on symbol and timeframe"""
    base_url = "https://stockcharts.com/c-sc/sc?chart="
    
    # Random parameter for cache busting
    import random
    rand_param = random.random()
    
    if timeframe == "5-Minute":
        return f"{base_url}{symbol},uu[900,700]iahayaca[bb][i!ub14!la12,26,9][pd20,2]&r={rand_param}"
    elif timeframe == "Daily":
        return f"{base_url}{symbol},uu[900,700]dahayaca[bb][i!ub14!lu][pd20,2]&r={rand_param}"
    elif timeframe == "Weekly":
        return f"{base_url}{symbol},uu[900,700]whhayaca[bb][i!ub14!lu][pd20,2]&r={rand_param}"
    return ""

def download_chart_image(url):
    """Download chart image from StockCharts and convert to base64"""
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.9',
            'Accept-Encoding': 'gzip, deflate, br',
            'Referer': 'https://stockcharts.com/',
            'Connection': 'keep-alive',
            'Sec-Fetch-Dest': 'image',
            'Sec-Fetch-Mode': 'no-cors',
            'Sec-Fetch-Site': 'same-origin'
        }
        
        response = requests.get(url, headers=headers, timeout=10)
        response.raise_for_status()
        
        # Convert to base64
        img_base64 = base64.b64encode(response.content).decode('utf-8')
        return f"data:image/png;base64,{img_base64}"
    except Exception as e:
        print(f"Error downloading image from {url}: {str(e)}")
        # Return a placeholder image URL or error message
        return None

def generate_single_stock_images(symbol):
    """Generate HTML for single stock analysis (3 charts)"""
    timeframes = ["5-Minute", "Daily", "Weekly"]
    
    html = f"<h2>股票分析: {symbol}</h2>"
    html += '<div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 20px; margin-top: 20px;">'
    
    for tf in timeframes:
        url = generate_chart_url(symbol, tf)
        img_data = download_chart_image(url)
        
        if img_data:
            html += f'''
            <div style="text-align: center;">
                <h3>{tf}</h3>
                <img src="{img_data}" alt="{symbol} - {tf}" style="max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px;">
            </div>
            '''
        else:
            html += f'''
            <div style="text-align: center;">
                <h3>{tf}</h3>
                <p style="color: red;">加载图片失败</p>
            </div>
            '''
    
    html += '</div>'
    return html

def update_single_stock(symbol):
    """Update single stock tab with charts"""
    if not symbol or not symbol.strip():
        return "<p>请输入股票代码</p>"
    
    symbol = symbol.strip().upper()
    return generate_single_stock_images(symbol)

def generate_watchlist_images(watchlist_name, timeframe):
    """Generate HTML for watchlist with selected timeframe"""
    if not watchlist_name or watchlist_name == "Select a watchlist":
        return "<p>请选择一个观察列表</p>"
    
    # Load watchlist config
    config = load_watchlists()
    if watchlist_name not in config:
        return f"<p>观察列表 '{watchlist_name}' 未找到</p>"
    
    # Load stocks from file
    filename = config[watchlist_name]
    stocks = load_stocks_from_file(filename)
    
    if not stocks:
        return f"<p>文件 {filename} 中没有找到股票</p>"
    
    # Generate HTML
    html = f"<h2>观察列表: {watchlist_name} ({timeframe})</h2>"
    html += '<div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 20px; margin-top: 20px;">'
    
    for symbol in stocks:
        url = generate_chart_url(symbol, timeframe)
        img_data = download_chart_image(url)
        
        if img_data:
            html += f'''
            <div style="text-align: center;">
                <h3>{symbol}</h3>
                <img src="{img_data}" alt="{symbol} - {timeframe}" style="max-width: 100%; height: auto; border: 1px solid #ddd; border-radius: 4px;">
            </div>
            '''
        else:
            html += f'''
            <div style="text-align: center;">
                <h3>{symbol}</h3>
                <p style="color: red;">加载图片失败</p>
            </div>
            '''
    
    html += '</div>'
    return html

def get_watchlist_names():
    """Get list of watchlist names from config"""
    config = load_watchlists()
    return ["Select a watchlist"] + list(config.keys())

def save_watchlist_config(config):
    """Save watchlist configuration to JSON file"""
    config_path = Path("watchlists_config.json")
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)

def save_stocks_to_file(filename, stocks_text):
    """Save stock symbols to a watchlist file"""
    file_path = Path(filename)
    with open(file_path, 'w') as f:
        f.write(stocks_text.strip())

def add_or_update_watchlist(watchlist_name, stocks_text):
    """Add a new watchlist or update an existing one"""
    if not watchlist_name or not watchlist_name.strip():
        return "❌ 请输入观察列表名称", gr.update(), gr.update()
    
    watchlist_name = watchlist_name.strip()
    
    if not stocks_text or not stocks_text.strip():
        return "❌ 请输入股票代码", gr.update(), gr.update()
    
    # Generate filename from watchlist name
    filename = watchlist_name.lower().replace(' ', '_') + '.txt'
    
    # Save stocks to file
    save_stocks_to_file(filename, stocks_text)
    
    # Update config
    config = load_watchlists()
    config[watchlist_name] = filename
    save_watchlist_config(config)
    
    # Get updated list for dropdown
    updated_choices = get_watchlist_names()
    
    return (
        f"✅ 观察列表 '{watchlist_name}' 已保存！",
        gr.update(choices=updated_choices),
        gr.update(choices=list(config.keys()), value=None)
    )

def delete_watchlist(watchlist_name):
    """Delete a watchlist"""
    if not watchlist_name:
        return "❌ 请选择要删除的观察列表", gr.update(), gr.update()
    
    config = load_watchlists()
    
    if watchlist_name not in config:
        return f"❌ 观察列表 '{watchlist_name}' 不存在", gr.update(), gr.update()
    
    # Get filename and delete file
    filename = config[watchlist_name]
    file_path = Path(filename)
    if file_path.exists():
        file_path.unlink()
    
    # Remove from config
    del config[watchlist_name]
    save_watchlist_config(config)
    
    # Get updated list for dropdown
    updated_choices = get_watchlist_names()
    
    return (
        f"✅ 观察列表 '{watchlist_name}' 已删除！",
        gr.update(choices=updated_choices),
        gr.update(choices=list(config.keys()) if config else [], value=None)
    )

def load_watchlist_for_edit(watchlist_name):
    """Load watchlist content for editing"""
    if not watchlist_name:
        return "", ""
    
    config = load_watchlists()
    if watchlist_name not in config:
        return "", ""
    
    filename = config[watchlist_name]
    stocks = load_stocks_from_file(filename)
    stocks_text = ", ".join(stocks)
    
    return watchlist_name, stocks_text

# Create Gradio interface
with gr.Blocks(title="股票分析仪表板") as demo:
    gr.Markdown("# 股票分析仪表板")
    
    with gr.Tabs():
        # Tab 1: Single Stock Analysis
        with gr.Tab("单股分析"):
            gr.Markdown("### 输入股票代码查看图表")
            
            with gr.Row():
                stock_input = gr.Textbox(
                    label="股票代码",
                    placeholder="例如: QQQ, AAPL, SPY",
                    scale=3
                )
                analyze_btn = gr.Button("分析", scale=1)
            
            single_stock_output = gr.HTML()
            
            # Event handlers
            analyze_btn.click(
                fn=update_single_stock,
                inputs=[stock_input],
                outputs=[single_stock_output]
            )
            
            stock_input.submit(
                fn=update_single_stock,
                inputs=[stock_input],
                outputs=[single_stock_output]
            )
        
        # Tab 2: Watchlist
        with gr.Tab("观察列表"):
            gr.Markdown("### 选择观察列表和时间周期查看所有图表")
            
            with gr.Row():
                watchlist_dropdown = gr.Dropdown(
                    choices=get_watchlist_names(),
                    value="Select a watchlist",
                    label="选择观察列表",
                    scale=2
                )
                
                timeframe_dropdown = gr.Dropdown(
                    choices=["5-Minute", "Daily", "Weekly"],
                    value="Daily",
                    label="选择时间周期",
                    scale=1
                )
            
            watchlist_output = gr.HTML()
            
            # Event handlers - update on any dropdown change
            watchlist_dropdown.change(
                fn=generate_watchlist_images,
                inputs=[watchlist_dropdown, timeframe_dropdown],
                outputs=[watchlist_output]
            )
            
            timeframe_dropdown.change(
                fn=generate_watchlist_images,
                inputs=[watchlist_dropdown, timeframe_dropdown],
                outputs=[watchlist_output]
            )
        
        # Tab 3: Manage Watchlists
        with gr.Tab("管理观察列表"):
            gr.Markdown("### 添加、修改或删除观察列表")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### 添加/修改观察列表")
                    
                    manage_name_input = gr.Textbox(
                        label="观察列表名称",
                        placeholder="例如: Tech Stocks"
                    )
                    
                    manage_stocks_input = gr.TextArea(
                        label="股票代码（逗号分隔）",
                        placeholder="例如: AAPL, MSFT, GOOGL, NVDA",
                        lines=5
                    )
                    
                    with gr.Row():
                        save_btn = gr.Button("💾 保存", variant="primary")
                    
                    save_status = gr.Markdown()
                
                with gr.Column():
                    gr.Markdown("#### 编辑/删除现有列表")
                    
                    edit_dropdown = gr.Dropdown(
                        choices=list(load_watchlists().keys()),
                        label="选择要编辑的观察列表",
                        value=None
                    )
                    
                    with gr.Row():
                        load_btn = gr.Button("📝 加载到左侧编辑")
                        delete_btn = gr.Button("🗑️ 删除", variant="stop")
            
            # Event handlers for management tab
            save_btn.click(
                fn=add_or_update_watchlist,
                inputs=[manage_name_input, manage_stocks_input],
                outputs=[save_status, watchlist_dropdown, edit_dropdown]
            )
            
            delete_btn.click(
                fn=delete_watchlist,
                inputs=[edit_dropdown],
                outputs=[save_status, watchlist_dropdown, edit_dropdown]
            )
            
            load_btn.click(
                fn=load_watchlist_for_edit,
                inputs=[edit_dropdown],
                outputs=[manage_name_input, manage_stocks_input]
            )

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