import os
import subprocess
import pandas as pd
import plotly.express as px
import gradio as gr
from dotenv import load_dotenv
from util import css, js, Color
from accounts import Account
from database import read_log
from typing import List

# Load environment variables
load_dotenv(override=True)

# Config
RUN_EVERY_N_MINUTES = int(os.getenv("RUN_EVERY_N_MINUTES", "60"))
RUN_EVEN_WHEN_MARKET_IS_CLOSED = os.getenv("RUN_EVEN_WHEN_MARKET_IS_CLOSED", "false").strip().lower() == "true"
USE_MANY_MODELS = os.getenv("USE_MANY_MODELS", "false").strip().lower() == "true"

names = ["Warren", "George", "Ray", "Cathie"]
lastnames = ["Patience", "Bold", "Systematic", "Crypto"]

if USE_MANY_MODELS:
    model_names = [
        "gpt-4.1-mini",
        "deepseek-chat",
        "gemini-2.5-flash-preview-04-17",
        "grok-3-mini-beta",
    ]
    short_model_names = ["GPT 4.1 Mini", "DeepSeek V3", "Gemini 2.5 Flash", "Grok 3 Mini"]
else:
    model_names = ["gpt-4o-mini"] * 4
    short_model_names = ["GPT 4o mini"] * 4


# UI Wrapper

mapper = {
    "trace": Color.WHITE,
    "agent": Color.CYAN,
    "function": Color.GREEN,
    "generation": Color.YELLOW,
    "response": Color.MAGENTA,
    "account": Color.RED,
}


class Trader:
    def __init__(self, name: str, lastname: str, model_name: str):
        self.name = name
        self.lastname = lastname
        self.model_name = model_name
        self.account = Account.get(name)

    def reload(self):
        self.account = Account.get(self.name)

    def get_title(self):
        return f"<div style='text-align: center;font-size:34px;'>{self.name}<span style='color:#ccc;font-size:24px;'> ({self.model_name}) - {self.lastname}</span></div>"

    def get_strategy(self):
        return self.account.get_strategy()

    def get_portfolio_value_df(self):
        df = pd.DataFrame(self.account.portfolio_value_time_series, columns=["datetime", "value"])
        df["datetime"] = pd.to_datetime(df["datetime"])
        return df

    def get_portfolio_value_chart(self):
        df = self.get_portfolio_value_df()
        fig = px.line(df, x="datetime", y="value")
        margin = dict(l=40, r=20, t=20, b=40)
        fig.update_layout(height=300, margin=margin, xaxis_title=None, yaxis_title=None,
                          paper_bgcolor="#bbb", plot_bgcolor="#dde")
        fig.update_xaxes(tickformat="%m/%d", tickangle=45, tickfont=dict(size=8))
        fig.update_yaxes(tickfont=dict(size=8), tickformat=",.0f")
        return fig

    def get_holdings_df(self):
        holdings = self.account.get_holdings()
        if not holdings:
            return pd.DataFrame(columns=["Symbol", "Quantity"])
        return pd.DataFrame([{"Symbol": s, "Quantity": q} for s, q in holdings.items()])

    def get_transactions_df(self):
        transactions = self.account.list_transactions()
        if not transactions:
            return pd.DataFrame(columns=["Timestamp", "Symbol", "Quantity", "Price", "Rationale"])
        return pd.DataFrame(transactions)

    def get_portfolio_value(self):
        value = self.account.calculate_portfolio_value() or 0.0
        pnl = self.account.calculate_profit_loss(value) or 0.0
        color = "green" if pnl >= 0 else "red"
        emoji = "⬆" if pnl >= 0 else "⬇"
        return f"<div style='text-align: center;background-color:{color};'><span style='font-size:32px'>${value:,.0f}</span><span style='font-size:24px'>&nbsp;&nbsp;&nbsp;{emoji}&nbsp;${pnl:,.0f}</span></div>"

    def get_logs(self, previous=None):
        logs = read_log(self.name, last_n=13)
        response = "".join(
            f"<span style='color:{mapper.get(t, Color.WHITE).value}'>{ts} : [{t}] {msg}</span><br/>"
            for ts, t, msg in logs
        )
        response = f"<div style='height:250px; overflow-y:auto;'>{response}</div>"
        return response if response != previous else gr.update()


class TraderView:
    def __init__(self, trader: Trader):
        self.trader = trader

    def make_ui(self):
        with gr.Column():
            gr.HTML(self.trader.get_title())
            self.portfolio_value = gr.HTML(self.trader.get_portfolio_value)
            self.chart = gr.Plot(self.trader.get_portfolio_value_chart)
            self.log = gr.HTML(self.trader.get_logs)
            self.holdings_table = gr.Dataframe(
                value=self.trader.get_holdings_df,
                label="Holdings", headers=["Symbol", "Quantity"], max_height=300
            )
            self.transactions_table = gr.Dataframe(
                value=self.trader.get_transactions_df,
                label="Recent Transactions",
                headers=["Timestamp", "Symbol", "Quantity", "Price", "Rationale"],
                max_height=300
            )

        gr.Timer(value=120).tick(fn=self.refresh, inputs=[], outputs=[
            self.portfolio_value, self.chart, self.holdings_table, self.transactions_table
        ], show_progress="hidden", queue=False)

        gr.Timer(value=0.5).tick(fn=self.trader.get_logs, inputs=[self.log], outputs=[self.log],
                                 show_progress="hidden", queue=False)

    def refresh(self):
        self.trader.reload()
        return (
            self.trader.get_portfolio_value(),
            self.trader.get_portfolio_value_chart(),
            self.trader.get_holdings_df(),
            self.trader.get_transactions_df(),
        )


def create_ui():
    traders = [Trader(n, ln, sm) for n, ln, sm in zip(names, lastnames, short_model_names)]
    trader_views = [TraderView(t) for t in traders]
    with gr.Blocks(title="Traders", css=css, js=js) as ui:
        with gr.Row():
            for tv in trader_views:
                tv.make_ui()
    return ui


if __name__ == "__main__":
    # Launch trading_floor.py as a separate subprocess
    trading_floor_proc = subprocess.Popen(["python", "trading_floor.py"])

    try:
        ui = create_ui()
        ui.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
    finally:
        trading_floor_proc.terminate()
        trading_floor_proc.wait()