# TrumpSignalData.py

from QuantConnect.Data import SubscriptionDataSource
from QuantConnect.Python import PythonData
from datetime import datetime, timedelta

class TrumpSignalData(PythonData):
    """
    Custom data class to hold signals generated from news analysis.
    """

    # (Optional) Declare your custom fields for clarity
    target_asset: str
    signal_type: str
    sentiment: str

    def get_source(self, config, date, isLiveMode):
        """
        Specifies the location of the CSV file containing your signals.
        Lean will default to REMOTE_FILE transport.
        """
        url = (
            "https://www.dropbox.com/scl/fi/4bwy13i1i8tpwsvs6jm2h/"
            "signals.csv?rlkey=ihs42wvkl0uy5u4w1tmfy0hkb&st=jn9wdib3&dl=1"
        )
        return SubscriptionDataSource(url)

    def reader(self, config, line, date, isLiveMode):
        """
        Parses one line of the CSV and returns a populated TrumpSignalData instance.
        """
        line = line.strip()
        # Skip empty lines and header rows
        if not line or line.lower().startswith("timestamp"):
            return None

        parts = line.split(',')
        time = datetime.strptime(parts[0], "%Y%m%d %H:%M:%S")

        # Create a fresh instance each time
        signal = TrumpSignalData()
        signal.Symbol    = config.Symbol
        signal.Time      = time
        signal.EndTime   = time + timedelta(minutes=1)

        # Extract your custom fields
        asset     = parts[1].strip()
        sig_type  = parts[2].strip()
        sentiment = parts[3].strip().lower()

        signal.target_asset = asset
        signal.signal_type  = sig_type
        signal.sentiment    = sentiment

        # Assign numeric Value based on sentiment
        if sentiment == "positive":
            signal.Value = 1.0
        elif sentiment == "negative":
            signal.Value = -1.0
        else:
            signal.Value = 0.0

        return signal

# main.py
from QuantConnect.Algorithm import QCAlgorithm
from TrumpSignalData import TrumpSignalData
from QuantConnect import Resolution

class TrumpTradingAgent(QCAlgorithm):

    def Initialize(self):
        self.set_start_date(2025, 7, 1)
        self.set_end_date(2025, 7, 3)
        self.set_cash(100_000)

        # Now works, because PythonData is instantiable
        self.add_data(TrumpSignalData, "TRUMPSIGNAL")

        self.traded_assets = set()
        self.add_equity("SPY", Resolution.MINUTE)
        self.add_equity("GLD", Resolution.MINUTE)

    def OnData(self, data):
        if not data.ContainsKey("TRUMPSIGNAL"):
            return

        signal = data["TRUMPSIGNAL"]
        asset     = signal.target_asset
        sig_type  = signal.signal_type
        sentiment = signal.sentiment

        self.log(f"Signal Received: Type='{sig_type}', Asset='{asset}', Sentiment='{sentiment}'")

        if asset != "NONE":
            if asset not in self.traded_assets:
                self.add_equity(asset, Resolution.MINUTE)
                self.traded_assets.add(asset)
                return

            if sig_type == "Bilateral_Deal_Boost" and sentiment == "positive":
                self.set_holdings(asset, 0.20)

            elif sig_type == "Punishment_Tweet" and sentiment == "negative":
                self.liquidate(asset)