from datetime import datetime, timedelta
from dataclasses import dataclass
from typing import Optional
import pytz
from zoneinfo import ZoneInfo
import yfinance as yf
from db.db_connection import create_client
from trading.position_calculator import PositionCalculator
from utils.data_utils import get_user_input

def handle_sell_order(ticker: str, shares_to_sell: int, exit_price: float, exit_date: datetime,
                     order_type: str, followed_rules: bool, exit_reason: str, notes: Optional[str] = None) -> bool:
    """
    Handle sell order using FIFO logic
    
    Args:
        ticker (str): Stock ticker
        shares_to_sell (int): Number of shares to sell
        exit_price (float): Exit price per share
        exit_date (datetime): Exit date and time
        order_type (str): Order type (Market/Limit)
        followed_rules (bool): Whether trading rules were followed
        exit_reason (str): Reason for exit
        notes (Optional[str]): Additional notes
        
    Returns:
        bool: True if sell order was processed successfully
    """
    with create_client() as client:
        # Get open positions for this ticker ordered by entry date (FIFO)
        query = f"""
            SELECT id, shares, entry_price, position_id
            FROM stock_db.trades
            WHERE ticker = '{ticker}'
            AND exit_price IS NULL
            ORDER BY entry_date ASC
        """
        result = client.query(query).result_rows
        
        if not result:
            print(f"No open positions found for {ticker}")
            return False
            
        remaining_shares = shares_to_sell
        positions = [dict(zip(['id', 'shares', 'entry_price', 'position_id'], row)) for row in result]
        total_available_shares = sum(pos['shares'] for pos in positions)
        
        if shares_to_sell > total_available_shares:
            print(f"Error: Attempting to sell {shares_to_sell} shares but only {total_available_shares} available")
            return False
            
        for position in positions:
            if remaining_shares <= 0:
                break
                
            shares_from_position = min(remaining_shares, position['shares'])
            
            if shares_from_position == position['shares']:
                # Close entire position
                update_trade(
                    trade_id=position['id'],
                    updates={
                        'exit_price': exit_price,
                        'exit_date': exit_date,
                        'followed_rules': 1 if followed_rules else 0,
                        'exit_reason': exit_reason,
                        'notes': notes
                    }
                )
            else:
                # Split position: update original position with remaining shares
                # and create a new closed position for the sold shares
                new_position_shares = position['shares'] - shares_from_position
                
                # Update original position with reduced shares
                client.command(f"""
                    ALTER TABLE stock_db.trades
                    UPDATE shares = {new_position_shares}
                    WHERE id = {position['id']}
                """)
                
                # Create new record for the sold portion
                trade = TradeEntry(
                    ticker=ticker,
                    entry_date=datetime.now(),  # Use original entry date?
                    shares=shares_from_position,
                    entry_price=position['entry_price'],
                    target_price=0,  # Not relevant for closed portion
                    stop_loss=0,  # Not relevant for closed portion
                    strategy="FIFO_SPLIT",
                    order_type=order_type,
                    position_id=position['position_id'],
                    followed_rules=followed_rules,
                    exit_price=exit_price,
                    exit_date=exit_date,
                    exit_reason=exit_reason,
                    notes=notes
                )
                add_trade(trade)
            
            remaining_shares -= shares_from_position
            
        return True

def create_portfolio_table():
    with create_client() as client:
        query = """
            CREATE TABLE IF NOT EXISTS stock_db.portfolio_history (
                id UInt32,
                date DateTime,
                total_value Float64,
                cash_balance Float64,
                notes Nullable(String),
                created_at DateTime DEFAULT now()
            ) ENGINE = MergeTree()
            ORDER BY (date, id)
        """
        client.command(query)

def update_portfolio_value(total_value: float, cash_balance: float, notes: Optional[str] = None):
    with create_client() as client:
        query = f"""
            INSERT INTO stock_db.portfolio_history (
                id, date, total_value, cash_balance, notes
            ) VALUES (
                {generate_id()},
                '{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}',
                {total_value},
                {cash_balance},
                {f"'{notes}'" if notes else 'NULL'}
            )
        """
        client.command(query)

def get_latest_portfolio_value() -> Optional[dict]:
    with create_client() as client:
        query = """
            SELECT total_value, cash_balance, date
            FROM stock_db.portfolio_history
            ORDER BY date DESC
            LIMIT 1
        """
        result = client.query(query).result_rows
        if result:
            return {
                'total_value': result[0][0],
                'cash_balance': result[0][1],
                'date': result[0][2]
            }
        return None

@dataclass
class TradeEntry:
    ticker: str
    entry_date: datetime
    shares: int
    entry_price: float
    target_price: Optional[float]  # Made optional since sell orders don't need it
    stop_loss: Optional[float]     # Made optional since sell orders don't need it
    strategy: Optional[str]        # Made optional since sell orders might not need it
    order_type: str               # Market/Limit
    direction: str                # 'buy' or 'sell'
    followed_rules: Optional[bool] = None
    entry_reason: Optional[str] = None
    exit_price: Optional[float] = None
    exit_date: Optional[datetime] = None
    exit_reason: Optional[str] = None
    notes: Optional[str] = None
    position_id: Optional[str] = None

    def __post_init__(self):
        # For sell orders, set exit_price and exit_date
        if self.direction.lower() == 'sell':
            self.exit_price = self.entry_price
            self.exit_date = self.entry_date
            self.entry_price = None  # Set to NULL for sells
            self.entry_date = None   # Set to NULL for sells

    @property
    def expected_profit_loss(self) -> Optional[float]:
        if self.direction == 'buy' and self.target_price:
            return (self.target_price - self.entry_price) * self.shares
        return None

    @property
    def max_loss(self) -> Optional[float]:
        if self.direction == 'buy' and self.stop_loss:
            return (self.stop_loss - self.entry_price) * self.shares
        return None
    
    @property
    def actual_profit_loss(self) -> Optional[float]:
        if self.exit_price:
            return (self.exit_price - self.entry_price) * self.shares
        return None

def get_market_hours(date: datetime) -> tuple:
    """Get market open/close times in Eastern for given date"""
    eastern = pytz.timezone('US/Eastern')
    date_eastern = date.astimezone(eastern)
    
    market_open = eastern.localize(
        datetime.combine(date_eastern.date(), datetime.strptime("09:30", "%H:%M").time())
    )
    market_close = eastern.localize(
        datetime.combine(date_eastern.date(), datetime.strptime("16:00", "%H:%M").time())
    )
    return market_open, market_close

def validate_market_time(dt: datetime) -> tuple[datetime, bool]:
    """
    Validate if time is during market hours, adjust if needed
    Returns: (adjusted_datetime, was_adjusted)
    """
    pacific = pytz.timezone('US/Pacific')
    eastern = pytz.timezone('US/Eastern')
    
    # Ensure datetime is timezone-aware
    if dt.tzinfo is None:
        dt = pacific.localize(dt)
    
    dt_eastern = dt.astimezone(eastern)
    market_open, market_close = get_market_hours(dt_eastern)
    
    if dt_eastern < market_open:
        return market_open.astimezone(pacific), True
    elif dt_eastern > market_close:
        return market_close.astimezone(pacific), True
    
    return dt, False

def get_datetime_input(prompt: str, default: datetime = None) -> Optional[datetime]:
    """Get date and time input in Pacific time"""
    pacific = pytz.timezone('US/Pacific')
    
    while True:
        try:
            if default:
                print(f"Press Enter for current time ({default.strftime('%Y-%m-%d %H:%M')})")
            date_str = input(f"{prompt} (YYYY-MM-DD HH:MM, q to quit): ").strip()
            
            if date_str.lower() in ['q', 'quit', 'exit']:
                return None
            
            if not date_str and default:
                dt = default
            else:
                dt = datetime.strptime(date_str, "%Y-%m-%d %H:%M")
            
            # Make datetime timezone-aware (Pacific)
            dt = pacific.localize(dt)
            
            # Validate market hours
            adjusted_dt, was_adjusted = validate_market_time(dt)
            if was_adjusted:
                print(f"\nWarning: Time adjusted to market hours (Eastern)")
                print(f"Original (Pacific): {dt.strftime('%Y-%m-%d %H:%M %Z')}")
                print(f"Adjusted (Pacific): {adjusted_dt.strftime('%Y-%m-%d %H:%M %Z')}")
                print(f"Adjusted (Eastern): {adjusted_dt.astimezone(pytz.timezone('US/Eastern')).strftime('%Y-%m-%d %H:%M %Z')}")
                if input("Accept adjusted time? (y/n): ").lower() != 'y':
                    continue
            
            return adjusted_dt
            
        except ValueError:
            print("Invalid format. Please use YYYY-MM-D HH:MM")

def create_trades_table():
    with create_client() as client:
        query = """
            CREATE TABLE IF NOT EXISTS stock_db.trades (
                id UInt32,
                position_id String,
                ticker String,
                entry_date DateTime,
                shares UInt32,
                entry_price Float64,
                target_price Nullable(Float64),
                stop_loss Nullable(Float64),
                strategy Nullable(String),
                order_type String,
                direction String,
                followed_rules Nullable(UInt8),
                entry_reason Nullable(String),
                exit_price Nullable(Float64),
                exit_date Nullable(DateTime),
                exit_reason Nullable(String),
                notes Nullable(String),
                created_at DateTime DEFAULT now(),
                plan_id Nullable(UInt32)
            ) ENGINE = MergeTree()
            ORDER BY (position_id, id, entry_date)
        """
        client.command(query)

def generate_id() -> int:
    """Generate a unique ID for the trade"""
    return int(datetime.now().timestamp() * 1000)

def generate_position_id(ticker: str, entry_date: datetime = None) -> str:
    """
    Generate a unique position ID for grouping related trades
    
    Args:
        ticker (str): Stock ticker symbol
        entry_date (datetime, optional): Entry date for the trade
        
    Returns:
        str: Position ID in format TICKER_YYYYMMDDHHMMSS
    """
    if entry_date is None:
        entry_date = datetime.now()
    timestamp = entry_date.strftime("%Y%m%d%H%M%S")
    return f"{ticker}_{timestamp}"

def get_position_summary(ticker: str) -> dict:
    """Get summary of existing positions for a ticker"""
    print(f"\n=== Getting Position Summary for {ticker} ===")  # Debug
    with create_client() as client:
        query = f"""
            SELECT 
                position_id,
                sum(shares) as total_shares,
                avg(entry_price) as avg_entry_price,
                min(entry_date) as first_entry,
                max(entry_date) as last_entry,
                count() as num_orders,
                any(target_price) as target_price,
                any(stop_loss) as stop_loss,
                any(strategy) as strategy
            FROM stock_db.trades 
            WHERE ticker = '{ticker}' 
            AND exit_price IS NULL 
            GROUP BY position_id
            ORDER BY first_entry DESC
        """
        print(f"Executing query: {query}")  # Debug
        result = client.query(query).result_rows
        print(f"Query returned {len(result)} rows")  # Debug
        
        columns = ['position_id', 'total_shares', 'avg_entry_price', 
                  'first_entry', 'last_entry', 'num_orders',
                  'target_price', 'stop_loss', 'strategy']
        positions = [dict(zip(columns, row)) for row in result]
        print(f"Processed positions: {positions}")  # Debug
        return positions

def get_order_type() -> Optional[str]:
    """Get order type from user"""
    while True:
        print("\nOrder Type:")
        print("1. Market")
        print("2. Limit")
        print("q. Quit")
        choice = input("Select order type (1-2, q to quit): ")
        
        if choice.lower() in ['q', 'quit', 'exit']:
            return None
        elif choice == "1":
            return "Market"
        elif choice == "2":
            return "Limit"
        print("Invalid choice. Please try again.")

def add_trade(trade: TradeEntry):
    """Add a new trade to the database"""
    with create_client() as client:
        # For sell orders, use exit_price and exit_date instead of entry_price and entry_date
        query = f"""
            INSERT INTO stock_db.trades (
                id, position_id, ticker, entry_date, shares, entry_price, target_price, 
                stop_loss, strategy, order_type, direction, followed_rules, entry_reason,
                exit_price, exit_date, exit_reason, notes
            ) VALUES (
                {generate_id()},
                '{trade.position_id}',
                '{trade.ticker}',
                {f"'{trade.entry_date.strftime('%Y-%m-%d %H:%M:%S')}'" if trade.entry_date else 'NULL'},
                {trade.shares},
                {trade.entry_price if trade.entry_price else 'NULL'},
                {trade.target_price if trade.target_price else 'NULL'},
                {trade.stop_loss if trade.stop_loss else 'NULL'},
                {f"'{trade.strategy}'" if trade.strategy else 'NULL'},
                '{trade.order_type}',
                '{trade.direction.lower()}',
                {1 if trade.followed_rules else 0},
                {f"'{trade.entry_reason}'" if trade.entry_reason else 'NULL'},
                {trade.exit_price if trade.exit_price else 'NULL'},
                {f"'{trade.exit_date.strftime('%Y-%m-%d %H:%M:%S')}'" if trade.exit_date else 'NULL'},
                {f"'{trade.exit_reason}'" if trade.exit_reason else 'NULL'},
                {f"'{trade.notes}'" if trade.notes else 'NULL'}
            )
        """
        client.command(query)

def update_trade(trade_id: int, updates: dict):
    """
    Update trade details
    
    Args:
        trade_id (int): ID of trade to update
        updates (dict): Dictionary of fields and values to update
    """
    with create_client() as client:
        # If trying to update entry_date, we need to delete and reinsert
        if 'entry_date' in updates:
            # First get the full trade data
            query = f"SELECT * FROM stock_db.trades WHERE id = {trade_id}"
            result = client.query(query).result_rows
            if not result:
                raise Exception("Trade not found")
                
            # Delete the existing trade
            client.command(f"ALTER TABLE stock_db.trades DELETE WHERE id = {trade_id}")
            
            # Prepare the new trade data
            columns = ['id', 'position_id', 'ticker', 'entry_date', 'shares', 'entry_price', 
                      'target_price', 'stop_loss', 'strategy', 'order_type', 'followed_rules', 
                      'entry_reason', 'exit_price', 'exit_date', 'exit_reason', 'notes', 'created_at']
            trade_data = dict(zip(columns, result[0]))
            trade_data.update(updates)
            
            # Insert the updated trade
            query = f"""
                INSERT INTO stock_db.trades (
                    id, position_id, ticker, entry_date, shares, entry_price, target_price, 
                    stop_loss, strategy, order_type, followed_rules, entry_reason, exit_price, 
                    exit_date, exit_reason, notes
                ) VALUES (
                    {trade_id},
                    '{trade_data['position_id']}',
                    '{trade_data['ticker']}',
                    '{trade_data['entry_date'].strftime('%Y-%m-%d %H:%M:%S')}',
                    {trade_data['shares']},
                    {trade_data['entry_price']},
                    {trade_data['target_price']},
                    {trade_data['stop_loss']},
                    '{trade_data['strategy']}',
                    '{trade_data['order_type']}',
                    {1 if trade_data['followed_rules'] else 0},
                    {f"'{trade_data['entry_reason']}'" if trade_data['entry_reason'] else 'NULL'},
                    {trade_data['exit_price'] if trade_data['exit_price'] else 'NULL'},
                    {f"'{trade_data['exit_date'].strftime('%Y-%m-%d %H:%M:%S')}'" if trade_data['exit_date'] else 'NULL'},
                    {f"'{trade_data['exit_reason']}'" if trade_data['exit_reason'] else 'NULL'},
                    {f"'{trade_data['notes']}'" if trade_data['notes'] else 'NULL'}
                )
            """
            client.command(query)
        else:
            # For non-key columns, we can use regular UPDATE
            update_statements = []
            for field, value in updates.items():
                if isinstance(value, str):
                    update_statements.append(f"{field} = '{value}'")
                elif isinstance(value, datetime):
                    update_statements.append(f"{field} = '{value.strftime('%Y-%m-%d %H:%M:%S')}'")
                elif value is None:
                    update_statements.append(f"{field} = NULL")
                else:
                    update_statements.append(f"{field} = {value}")
                    
            update_clause = ", ".join(update_statements)
            
            query = f"""
                ALTER TABLE stock_db.trades
                UPDATE {update_clause}
                WHERE id = {trade_id}
            """
            client.command(query)

def get_open_trades_summary() -> dict:
    """Get summary of all open trades grouped by ticker"""
    print("\n=== Fetching Open Trades Summary ===")  # Debug
    with create_client() as client:
        query = """
            WITH position_totals AS (
                SELECT 
                    ticker,
                    position_id,
                    sum(CASE 
                        WHEN direction = 'sell' THEN -shares
                        ELSE shares 
                    END) as net_shares
                FROM stock_db.trades
                GROUP BY ticker, position_id
                HAVING net_shares > 0
            )
            SELECT 
                t.ticker,
                sum(CASE 
                    WHEN t.direction = 'sell' THEN -t.shares 
                    ELSE t.shares 
                END) as total_shares,
                avg(t.entry_price) as avg_entry_price,
                min(t.entry_date) as first_entry,
                max(t.entry_date) as last_entry,
                count() as num_orders,
                groupArray(t.position_id) as position_ids
            FROM stock_db.trades t
            INNER JOIN position_totals pt 
                ON t.ticker = pt.ticker 
                AND t.position_id = pt.position_id
            GROUP BY t.ticker
            HAVING total_shares > 0
            ORDER BY t.ticker ASC
        """
        print(f"Executing summary query: {query}")  # Debug
        try:
            result = client.query(query).result_rows
            print(f"Summary query returned {len(result)} rows")  # Debug
            
            columns = ['ticker', 'total_shares', 'avg_entry_price', 
                      'first_entry', 'last_entry', 'num_orders', 'position_ids']
            summaries = [dict(zip(columns, row)) for row in result]
            print(f"Processed summaries: {summaries}")  # Debug
            return summaries
        except Exception as e:
            print(f"Error in get_open_trades_summary: {str(e)}")  # Debug
            raise

def get_open_trades():
    print("\n=== Fetching Open Trades ===")  # Debug
    with create_client() as client:
        query = """
            WITH position_totals AS (
                SELECT 
                    ticker,
                    position_id,
                    sum(CASE 
                        WHEN direction = 'sell' THEN -shares
                        ELSE shares 
                    END) as net_shares
                FROM stock_db.trades
                GROUP BY ticker, position_id
                HAVING net_shares > 0
            )
            SELECT t.* 
            FROM stock_db.trades t
            INNER JOIN position_totals pt 
                ON t.ticker = pt.ticker 
                AND t.position_id = pt.position_id
            ORDER BY t.entry_date DESC
        """
        print(f"Executing query: {query}")  # Debug
        try:
            result = client.query(query).result_rows
            print(f"Query returned {len(result)} rows")  # Debug
            
            columns = ['id', 'position_id', 'ticker', 'entry_date', 'shares', 'entry_price', 'target_price', 
                      'stop_loss', 'strategy', 'order_type', 'followed_rules', 'entry_reason', 'exit_price', 
                      'exit_date', 'exit_reason', 'notes', 'created_at']
            trades = [dict(zip(columns, row)) for row in result]
            print(f"Processed trades: {trades}")  # Debug
            return trades
        except Exception as e:
            print(f"Error in get_open_trades: {str(e)}")  # Debug
            raise

def get_current_prices(tickers: list) -> dict:
    """Get current prices for multiple tickers using yfinance"""
    prices = {}
    for ticker in tickers:
        try:
            stock = yf.Ticker(ticker)
            prices[ticker] = stock.fast_info["last_price"]
        except Exception as e:
            print(f"Error getting price for {ticker}: {e}")
    return prices

def delete_trade(trade_id: int) -> bool:
    """
    Delete a trade from the database
    
    Args:
        trade_id (int): ID of trade to delete
        
    Returns:
        bool: True if deletion was successful
    """
    try:
        with create_client() as client:
            query = f"""
                ALTER TABLE stock_db.trades
                DELETE WHERE id = {trade_id}
            """
            client.command(query)
            return True
    except Exception as e:
        print(f"Error deleting trade: {e}")
        return False

def get_trade_history(limit: int = 50):
    with create_client() as client:
        query = f"""
            SELECT 
                id, position_id, ticker, entry_date, shares, entry_price, 
                target_price, stop_loss, strategy, order_type, direction,
                followed_rules, entry_reason, exit_price, exit_date, 
                exit_reason, notes, created_at,
                groupArray(ticker) OVER (PARTITION BY position_id) as related_tickers,
                groupArray(id) OVER (PARTITION BY position_id) as related_ids
            FROM stock_db.trades 
            ORDER BY entry_date DESC 
            LIMIT {limit}
        """
        result = client.query(query).result_rows
        columns = ['id', 'position_id', 'ticker', 'entry_date', 'shares', 
                  'entry_price', 'target_price', 'stop_loss', 'strategy', 
                  'order_type', 'direction', 'followed_rules', 'entry_reason', 
                  'exit_price', 'exit_date', 'exit_reason', 'notes', 
                  'created_at', 'related_tickers', 'related_ids']
        return [dict(zip(columns, row)) for row in result]

def journal_menu():
    """Trading journal menu interface"""
    create_trades_table()  # Ensure table exists
    
    while True:
        print("\nTrading Journal")
        print("1. Add New Trade")
        print("2. Update Existing Trade")
        print("3. View Open Trades")
        print("4. View Trade History")
        print("5. Delete Trade")
        print("6. Return to Main Menu")
        
        choice = input("\nSelect an option (1-5): ")
        
        if choice == "1":
            ticker = get_user_input("Enter ticker symbol:", str)
            if ticker is None:
                continue
            ticker = ticker.upper()
            
            # Ask if this is a buy or sell order
            print("\nOrder Direction:")
            print("1. Buy")
            print("2. Sell")
            direction = get_user_input("Select direction (1-2):", str)
            if direction is None:
                continue
                
            if direction not in ["1", "2"]:
                print("Invalid direction")
                continue
                
            if direction == "2":  # Sell order
                shares = get_user_input("Enter number of shares to sell:", int)
                if shares is None:
                    continue
                    
                exit_price = get_user_input("Enter exit price:", float)
                if exit_price is None:
                    continue
                    
                exit_date = get_datetime_input("Enter exit date and time", default=datetime.now())
                if exit_date is None:
                    continue
                    
                order_type = get_order_type()
                if order_type is None:
                    continue
                    
                followed_rules = get_user_input("Did you follow your rules? (y/n):", bool)
                if followed_rules is None:
                    continue
                    
                exit_reason = input("Enter exit reason: ")
                notes = input("Additional notes (optional): ") or None
                
                if handle_sell_order(
                    ticker=ticker,
                    shares_to_sell=shares,
                    exit_price=exit_price,
                    exit_date=exit_date,
                    order_type=order_type,
                    followed_rules=followed_rules,
                    exit_reason=exit_reason,
                    notes=notes
                ):
                    print("Sell order processed successfully!")
                continue
            
            # Show existing positions for this ticker
            existing_positions = get_position_summary(ticker)
            if existing_positions:
                print(f"\nExisting {ticker} Positions:")
                for pos in existing_positions:
                    print(f"\nPosition ID: {pos['position_id']}")
                    print(f"Total Shares: {pos['total_shares']}")
                    print(f"Average Entry: ${pos['avg_entry_price']:.2f}")
                    print(f"First Entry: {pos['first_entry']}")
                    print(f"Number of Orders: {pos['num_orders']}")
                
                add_to_existing = get_user_input("Add to existing position? (y/n):", bool)
                if add_to_existing is None:
                    continue
                    
                if add_to_existing:
                    position_id = get_user_input("Enter Position ID:", str)
                    if position_id is None:
                        continue
                else:
                    position_id = generate_position_id(ticker)
            else:
                position_id = generate_position_id(ticker)
            
            # Get entry date/time with market hours validation
            entry_date = get_datetime_input("Enter entry date and time", default=datetime.now())
            if entry_date is None:
                continue
            
            shares = get_user_input("Enter number of shares:", int)
            if shares is None:
                continue
                
            entry_price = get_user_input("Enter entry price:", float)
            if entry_price is None:
                continue
                
            order_type = get_order_type()
            if order_type is None:
                continue
            
            # If adding to existing position, get target/stop from existing
            if existing_positions and add_to_existing:
                # Use existing target and stop loss
                target_price = float(input("Enter new target price (or press Enter to keep existing): ") or 
                                  existing_positions[0]['target_price'])
                stop_loss = float(input("Enter new stop loss (or press Enter to keep existing): ") or 
                                existing_positions[0]['stop_loss'])
                strategy = input("Enter strategy name (or press Enter to keep existing): ") or existing_positions[0]['strategy']
            else:
                target_price = float(input("Enter target price: "))
                stop_loss = float(input("Enter stop loss: "))
                strategy = input("Enter strategy name: ")
            
            followed_rules = input("Did you follow your rules? (y/n): ").lower() == 'y'
            entry_reason = input("Enter entry reason (optional): ") or None
            notes = input("Additional notes (optional): ") or None
            
            trade = TradeEntry(
                ticker=ticker,
                entry_date=entry_date,
                shares=shares,
                entry_price=entry_price,
                target_price=target_price,
                stop_loss=stop_loss,
                strategy=strategy,
                order_type=order_type,
                position_id=position_id,
                followed_rules=followed_rules,
                entry_reason=entry_reason,
                notes=notes
            )
            
            add_trade(trade)
            
            # Show updated position summary
            updated_positions = get_position_summary(ticker)
            if updated_positions:
                pos = updated_positions[0]  # Get the most recent position
                print(f"\nUpdated Position Summary for {ticker}:")
                print(f"Total Shares: {pos['total_shares']}")
                print(f"Average Entry: ${pos['avg_entry_price']:.2f}")
                print(f"Expected Profit: ${(target_price - pos['avg_entry_price']) * pos['total_shares']:.2f}")
                print(f"Maximum Loss: ${(stop_loss - pos['avg_entry_price']) * pos['total_shares']:.2f}")
            
            print("Trade added successfully!")
            
        elif choice == "2":
            open_trades = get_open_trades()
            if not open_trades:
                print("No open trades to update.")
                continue
                
            print("\nOpen Trades:")
            for trade in open_trades:
                print(f"{trade['id']}: {trade['ticker']} - Entered at ${trade['entry_price']}")
            
            print("\nOpen Trades:")
            for trade in open_trades:
                print(f"\nID: {trade['id']}")
                print(f"Ticker: {trade['ticker']}")
                print(f"Entry Date: {trade['entry_date']}")
                print(f"Shares: {trade['shares']}")
                print(f"Entry Price: ${trade['entry_price']}")
                print(f"Target: ${trade['target_price']}")
                print(f"Stop Loss: ${trade['stop_loss']}")
                print(f"Strategy: {trade['strategy']}")
                print("-" * 40)
            
            trade_id = get_user_input("\nEnter trade ID to update:", int)
            if trade_id is None:
                continue
                
            # Find the trade to update
            trade_to_update = next((t for t in open_trades if t['id'] == trade_id), None)
            if not trade_to_update:
                print("Trade not found.")
                continue
            
            print("\nUpdate Trade Fields")
            print("Leave blank to keep existing value")
            
            updates = {}
            
            # Entry date
            new_date = get_datetime_input("Enter new entry date and time (blank to keep):", None)
            if new_date:
                updates['entry_date'] = new_date
            
            # Shares
            new_shares = get_user_input("Enter new number of shares:", int, allow_empty=True)
            if new_shares is not None:
                updates['shares'] = new_shares
            
            # Entry price
            new_entry = get_user_input("Enter new entry price:", float, allow_empty=True)
            if new_entry is not None:
                updates['entry_price'] = new_entry
            
            # Target price
            new_target = get_user_input("Enter new target price:", float, allow_empty=True)
            if new_target is not None:
                updates['target_price'] = new_target
            
            # Stop loss
            new_stop = get_user_input("Enter new stop loss:", float, allow_empty=True)
            if new_stop is not None:
                updates['stop_loss'] = new_stop
            
            # Strategy
            new_strategy = input("Enter new strategy (blank to keep): ").strip()
            if new_strategy:
                updates['strategy'] = new_strategy
            
            # Order type
            if input("Update order type? (y/n): ").lower() == 'y':
                new_order_type = get_order_type()
                if new_order_type:
                    updates['order_type'] = new_order_type
            
            # Notes
            new_notes = input("Enter new notes (blank to keep): ").strip()
            if new_notes:
                updates['notes'] = new_notes
            
            if updates:
                try:
                    update_trade(trade_id, updates)
                    print("Trade updated successfully!")
                except Exception as e:
                    print(f"Error updating trade: {e}")
            else:
                print("No updates provided.")
            
        elif choice == "3":
            open_trades = get_open_trades()
            open_summary = get_open_trades_summary()
            
            if not open_trades:
                print("No open trades found.")
            else:
                # Get current prices for all open positions
                unique_tickers = list(set(summary['ticker'] for summary in open_summary))
                current_prices = get_current_prices(unique_tickers)
                
                print("\n=== Open Trades Summary ===")
                total_portfolio_value = 0
                total_paper_pl = 0
                
                for summary in open_summary:
                    ticker = summary['ticker']
                    avg_entry = summary['avg_entry_price']
                    current_price = current_prices.get(ticker)
                    
                    stop_loss = avg_entry * 0.93  # 7% stop loss
                    total_shares = summary['total_shares']
                    position_value = avg_entry * total_shares
                    max_loss = (avg_entry - stop_loss) * total_shares
                    
                    print(f"\n{ticker} Summary:")
                    print(f"Total Shares: {total_shares}")
                    print(f"Average Entry: ${avg_entry:.2f}")
                    print(f"Total Position Value: ${position_value:.2f}")
                    print(f"Combined Stop Loss (7%): ${stop_loss:.2f}")
                    print(f"Maximum Loss at Stop: ${max_loss:.2f}")
                    
                    if current_price:
                        current_value = current_price * total_shares
                        paper_pl = (current_price - avg_entry) * total_shares
                        pl_percentage = (paper_pl / position_value) * 100
                        total_portfolio_value += current_value
                        total_paper_pl += paper_pl
                        
                        print(f"Current Price: ${current_price:.2f}")
                        print(f"Current Value: ${current_value:.2f}")
                        print(f"Paper P/L: ${paper_pl:.2f} ({pl_percentage:.2f}%)")
                    
                    print(f"Number of Orders: {summary['num_orders']}")
                    print(f"Position Duration: {summary['last_entry'] - summary['first_entry']}")
                    print("-" * 50)
                
                if total_portfolio_value > 0:
                    print(f"\nTotal Portfolio Value: ${total_portfolio_value:.2f}")
                    print(f"Total Paper P/L: ${total_paper_pl:.2f}")
                    print(f"Overall P/L %: {(total_paper_pl / (total_portfolio_value - total_paper_pl)) * 100:.2f}%")
                
                print("\n=== Individual Trades ===")
                for trade in open_trades:
                    ticker = trade['ticker']
                    current_price = current_prices.get(ticker)
                    
                    print(f"\nTicker: {ticker}")
                    print(f"Position ID: {trade['position_id']}")
                    print(f"Entry Date: {trade['entry_date']}")
                    print(f"Shares: {trade['shares']}")
                    print(f"Entry Price: ${trade['entry_price']}")
                    print(f"Target: ${trade['target_price']}")
                    print(f"Stop Loss: ${trade['stop_loss']}")
                    print(f"Strategy: {trade['strategy']}")
                    print(f"Order Type: {trade['order_type']}")
                    
                    if current_price:
                        paper_pl = (current_price - trade['entry_price']) * trade['shares']
                        pl_percentage = (paper_pl / (trade['entry_price'] * trade['shares'])) * 100
                        print(f"Current Price: ${current_price:.2f}")
                        print(f"Paper P/L: ${paper_pl:.2f} ({pl_percentage:.2f}%)")
                    
                    if trade['entry_reason']:
                        print(f"Entry Reason: {trade['entry_reason']}")
                    if trade['notes']:
                        print(f"Notes: {trade['notes']}")
                    print("-" * 40)
                    
        elif choice == "4":
            history = get_trade_history()
            if not history:
                print("No trade history found.")
            else:
                print("\nTrade History:")
                for trade in history:
                    profit_loss = (trade['exit_price'] - trade['entry_price']) * trade['shares'] if trade['exit_price'] else None
                    print(f"\nTicker: {trade['ticker']}")
                    print(f"Entry: ${trade['entry_price']} on {trade['entry_date']}")
                    if trade['exit_price']:
                        print(f"Exit: ${trade['exit_price']} on {trade['exit_date']}")
                        print(f"P/L: ${profit_loss:.2f}")
                    print(f"Strategy: {trade['strategy']}")
                    if trade['notes']:
                        print(f"Notes: {trade['notes']}")
                    print("-" * 40)
                    
        elif choice == "5":
            # Show all trades (both open and closed)
            print("\nAll Trades:")
            with create_client() as client:
                query = "SELECT * FROM stock_db.trades ORDER BY entry_date DESC"
                result = client.query(query).result_rows
                columns = ['id', 'position_id', 'ticker', 'entry_date', 'shares', 'entry_price', 
                          'target_price', 'stop_loss', 'strategy', 'order_type', 'followed_rules', 
                          'entry_reason', 'exit_price', 'exit_date', 'exit_reason', 'notes', 'created_at']
                trades = [dict(zip(columns, row)) for row in result]
                
                for trade in trades:
                    print(f"\nID: {trade['id']}")
                    print(f"Ticker: {trade['ticker']}")
                    print(f"Entry Date: {trade['entry_date']}")
                    print(f"Shares: {trade['shares']}")
                    print(f"Entry Price: ${trade['entry_price']}")
                    if trade['exit_price']:
                        print(f"Exit Price: ${trade['exit_price']}")
                        print(f"Exit Date: {trade['exit_date']}")
                        pl = (trade['exit_price'] - trade['entry_price']) * trade['shares']
                        print(f"P/L: ${pl:.2f}")
                    print("-" * 40)
            
            trade_id = get_user_input("\nEnter trade ID to delete:", int)
            if trade_id is None:
                continue
                
            confirm = input(f"Are you sure you want to delete trade {trade_id}? (y/n): ").lower()
            if confirm == 'y':
                if delete_trade(trade_id):
                    print("Trade deleted successfully!")
                else:
                    print("Failed to delete trade.")
            else:
                print("Delete cancelled.")
                
        elif choice == "6":
            break