import gradio as gr
import pandas as pd
from futu import OpenQuoteContext
import threading
import time
import os
import json

PERMANENT_BLOCKLIST_FILE = "permanent_blocklist.json"

def load_permanent_blocklist():
    if os.path.exists(PERMANENT_BLOCKLIST_FILE):
        with open(PERMANENT_BLOCKLIST_FILE, 'r') as f:
            return set(json.load(f))
    return set()

def save_permanent_blocklist(blocklist):
    with open(PERMANENT_BLOCKLIST_FILE, 'w') as f:
        json.dump(list(blocklist), f)

def stock_filter(data, filters, perm_blocklist, sess_blocklist):
    symbol = data['code']
    if symbol in perm_blocklist or symbol in sess_blocklist:
        return False
    price = data['last_price']
    volume_1min = data['cur_vol']
    day_volume = data['turnover']
    spread = abs(data['ask_price'] - data['bid_price']) / ((data['ask_price'] + data['bid_price']) / 2) * 100 if (data['ask_price'] + data['bid_price']) > 0 else 0

    if not (filters['min_price'] <= price <= filters['max_price']):
        return False
    if volume_1min < filters['min_1min_vol']:
        return False
    if day_volume < filters['min_day_vol']:
        return False
    if spread > filters['max_spread']:
        return False
    return True

def poll_stocks(api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, 
                perm_blocklist_input, sess_blocklist_input, start_event, stop_event, history, table_callback):

    perm_blocklist = set(perm_blocklist_input.strip().upper().splitlines()) | load_permanent_blocklist()
    sess_blocklist = set(sess_blocklist_input.strip().upper().splitlines())
    filters = {
        "min_price": min_price,
        "max_price": max_price,
        "min_1min_vol": min_1min_vol,
        "min_day_vol": min_day_vol,
        "max_spread": max_spread,
    }
    save_permanent_blocklist(perm_blocklist)
    q = OpenQuoteContext(host='127.0.0.1', port=11111)  # Futu OpenD must be running

    # For demo: use a small static universe (replace with a live pull in production)
    codes = ['AAPL', 'MSFT', 'AMZN', 'GOOGL', 'TSLA']
    seen = set()
    while start_event.is_set() and not stop_event.is_set():
        for code in codes:
            ret, data = q.get_quote([code])
            if ret == 0 and not data.empty:
                record = data.iloc[0].to_dict()
                if stock_filter(record, filters, perm_blocklist, sess_blocklist):
                    stats = {
                        "Code": record["code"],
                        "Price": record["last_price"],
                        "1m Vol": record["cur_vol"],
                        "Day Vol": record["turnover"],
                        "Spread %": abs(record['ask_price'] - record['bid_price']) / ((record['ask_price'] + record['bid_price']) / 2) * 100
                    }
                    row_key = (stats['Code'], stats['Price'], stats['1m Vol'], stats['Day Vol'])
                    if row_key not in seen:
                        history.insert(0, stats)
                        seen.add(row_key)
                        table_callback(pd.DataFrame(history))
        time.sleep(1)
    q.close()

def start_trading(api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, 
                  perm_blocklist, sess_blocklist):
    start_event = threading.Event()
    stop_event = threading.Event()
    start_event.set()
    history = []
    table = pd.DataFrame()

    def table_callback(df):
        nonlocal table
        table = df

    thread = threading.Thread(target=poll_stocks, args=(
        api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, 
        perm_blocklist, sess_blocklist, start_event, stop_event, history, table_callback
    ))
    thread.start()

    def stop():
        stop_event.set()
        thread.join()
        csv = pd.DataFrame(history).to_csv(index=False)
        return table, csv

    return stop

with gr.Blocks() as demo:
    gr.Markdown("# MooMoo After Hours Stock Filter Tool")
    api_key = gr.Textbox(label="MooMoo API Key")
    min_price = gr.Number(label="Min Price", value=1)
    max_price = gr.Number(label="Max Price", value=20)
    min_1min_vol = gr.Number(label="Min 1-Minute Volume", value=5000)
    min_day_vol = gr.Number(label="Min Day Volume", value=100000)
    max_spread = gr.Number(label="Max Bid/Ask Spread %", value=4.5)
    perm_blocklist = gr.Textbox(label="Permanent Blocklist (1 per line)")
    sess_blocklist = gr.Textbox(label="Session Blocklist (1 per line)")
    table = gr.Dataframe(headers=["Code", "Price", "1m Vol", "Day Vol", "Spread %"])
    download = gr.File(label="Download Session Log")
    start_btn = gr.Button("Start")
    stop_btn = gr.Button("Stop")
    stop_fn_state = gr.State(None)

    def on_start(api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, 
                 perm_blocklist, sess_blocklist, _):
        stop_fn = start_trading(api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, perm_blocklist, sess_blocklist)
        return gr.update(interactive=False), gr.update(interactive=True), stop_fn

    def on_stop(stop_fn):
        if stop_fn:
            table, csv = stop_fn()
            fname = "session_log.csv"
            with open(fname, "w") as f:
                f.write(csv)
            return gr.update(interactive=True), gr.update(interactive=False), table, fname, None
        return gr.update(interactive=True), gr.update(interactive=False), None, None, None

    start_btn.click(
        on_start,
        inputs=[api_key, min_price, max_price, min_1min_vol, min_day_vol, max_spread, perm_blocklist, sess_blocklist, stop_fn_state],
        outputs=[start_btn, stop_btn, stop_fn_state]
    )

    stop_btn.click(
        on_stop,
        inputs=[stop_fn_state],
        outputs=[start_btn, stop_btn, table, download, stop_fn_state]
    )

demo.launch()