import streamlit as st
import plotly.graph_objects as go
from datetime import datetime
import pytz
from trading.journal import (
    create_trades_table, get_open_trades, get_trade_history,
    add_trade, update_trade, delete_trade, TradeEntry,
    get_open_trades_summary, get_current_prices, generate_position_id,
    get_position_summary, get_latest_portfolio_value, update_portfolio_value
)

def calculate_position_performance(trades):
    """Calculate performance metrics for a group of trades"""
    total_bought = 0
    total_cost = 0
    total_sold = 0
    total_proceeds = 0
    
    for trade in trades:
        try:
            shares = float(trade['shares'])
            
            # Handle buy orders
            if ('direction' not in trade or not trade['direction'] or 
                trade.get('direction', '').lower() == 'buy'):
                if trade.get('entry_price'):
                    price = float(trade['entry_price'])
                    total_bought += shares
                    total_cost += shares * price
            
            # Handle sell orders
            elif trade.get('direction', '').lower() == 'sell' or trade.get('exit_price'):
                price = float(trade.get('exit_price') or trade.get('entry_price'))
                total_sold += shares
                total_proceeds += shares * price
                
            print(f"Processing trade: Direction={trade.get('direction')}, "
                  f"Shares={shares}, Price={price}")  # Debug
            
        except (ValueError, TypeError) as e:
            print(f"Error processing trade: {e}")
            continue
    
    print(f"Performance calculation results:")  # Debug
    print(f"Total bought: {total_bought} shares at total cost: ${total_cost}")  # Debug
    print(f"Total sold: {total_sold} shares with total proceeds: ${total_proceeds}")  # Debug
    
    # Avoid division by zero
    if total_bought == 0:
        return {
            'total_bought': 0,
            'total_sold': 0,
            'avg_entry': 0,
            'avg_exit': 0,
            'realized_pl': 0,
            'remaining_shares': 0
        }
        
    avg_entry = total_cost / total_bought if total_bought > 0 else 0
    avg_exit = total_proceeds / total_sold if total_sold > 0 else 0
    realized_pl = total_proceeds - (total_sold / total_bought * total_cost) if total_sold > 0 else 0
    remaining_shares = total_bought - total_sold
    
    return {
        'total_bought': total_bought,
        'total_sold': total_sold,
        'avg_entry': avg_entry,
        'avg_exit': avg_exit,
        'realized_pl': realized_pl,
        'remaining_shares': remaining_shares
    }

def format_datetime(dt):
    """Format datetime for display"""
    if dt:
        return dt.strftime('%Y-%m-%d %H:%M')
    return ''

def plot_trade_history(trades):
    """Create a P/L chart using Plotly"""
    if not trades:
        return None
        
    # Prepare data
    dates = []
    pnl = []
    cumulative_pnl = 0
    
    for trade in trades:
        # Skip trades without numeric exit prices
        if not trade['exit_price'] or not isinstance(trade['exit_price'], (int, float)):
            continue
            
        try:
            # Convert prices to float if they're strings and are numeric
            exit_price = float(trade['exit_price'])
            entry_price = float(trade['entry_price'])
            shares = float(trade['shares'])
            
            trade_pnl = (exit_price - entry_price) * shares
            cumulative_pnl += trade_pnl
            dates.append(trade['exit_date'])
            pnl.append(cumulative_pnl)
        except (ValueError, TypeError):
            # Skip trades where conversion to float fails
            continue
    
    if not dates:
        return None
        
    # Create figure
    fig = go.Figure()
    fig.add_trace(
        go.Scatter(x=dates, y=pnl, mode='lines+markers',
                  name='Cumulative P/L',
                  line=dict(color='blue'),
                  hovertemplate='Date: %{x}<br>P/L: $%{y:.2f}<extra></extra>')
    )
    
    fig.update_layout(
        title='Cumulative Profit/Loss Over Time',
        xaxis_title='Date',
        yaxis_title='Cumulative P/L ($)',
        hovermode='x unified'
    )
    
    return fig

def trading_journal_page():
    st.header("Trading Journal")
    print("\n=== Starting Trading Journal Page ===")  # Debug
    
    # Tabs for different journal functions
    tab1, tab2, tab3, tab4 = st.tabs(["Open Positions", "Add Trade", "Update Trade", "Trade History"])
    
    with tab1:
        st.subheader("Open Positions")
        print("\n--- Fetching Open Positions ---")  # Debug
        open_trades = get_open_trades()
        print(f"Retrieved {len(open_trades) if open_trades else 0} open trades")  # Debug
        print(f"Open trades  {open_trades}")  # Debug
        
        open_summary = get_open_trades_summary()
        print(f"Retrieved {len(open_summary) if open_summary else 0} position summaries")  # Debug
        print(f"Open summary  {open_summary}")  # Debug
        
        if open_summary:
            # Get current prices
            unique_tickers = list(set(summary['ticker'] for summary in open_summary))
            print(f"Fetching prices for tickers: {unique_tickers}")  # Debug
            current_prices = get_current_prices(unique_tickers)
            print(f"Retrieved current prices: {current_prices}")  # Debug
            
            total_portfolio_value = 0
            total_paper_pl = 0
            
            for summary in open_summary:
                print(f"\nProcessing summary for {summary['ticker']}")  # Debug
                with st.expander(f"{summary['ticker']} Summary"):
                    ticker = summary['ticker']
                    avg_entry = summary['avg_entry_price']
                    current_price = current_prices.get(ticker)
                    total_shares = summary['total_shares']
                    position_value = avg_entry * total_shares
                    
                    print(f"Position details for {ticker}:")  # Debug
                    print(f"- Shares: {total_shares}")  # Debug
                    print(f"- Avg Entry: {avg_entry}")  # Debug
                    print(f"- Current Price: {current_price}")  # Debug
                    print(f"- Position Value: {position_value}")  # Debug
                    
                    col1, col2 = st.columns(2)
                    with col1:
                        st.metric("Total Shares", f"{total_shares:,}")
                        st.metric("Average Entry", f"${avg_entry:.2f}")
                        st.metric("Position Value", f"${position_value:.2f}")
                    
                    with col2:
                        if current_price:
                            current_value = current_price * total_shares
                            paper_pl = (current_price - avg_entry) * total_shares
                            pl_percentage = (paper_pl / position_value) * 100
                            
                            st.metric("Current Price", f"${current_price:.2f}")
                            st.metric("Paper P/L", f"${paper_pl:.2f}", f"{pl_percentage:.2f}%")
                            
                            total_portfolio_value += current_value
                            total_paper_pl += paper_pl
            
            if total_portfolio_value > 0:
                st.markdown("---")
                st.subheader("Portfolio Summary")
                col1, col2 = st.columns(2)
                with col1:
                    st.metric("Total Portfolio Value", f"${total_portfolio_value:.2f}")
                with col2:
                    st.metric("Total P/L", f"${total_paper_pl:.2f}", 
                             f"{(total_paper_pl / (total_portfolio_value - total_paper_pl)) * 100:.2f}%")
    
    with tab2:
        st.subheader("Add New Trade")
        print("\n--- Add Trade Tab ---")  # Debug
        
        ticker = st.text_input("Ticker Symbol").upper()
        print(f"Ticker entered: {ticker}")  # Debug
        
        # Add direction selection
        direction = st.selectbox(
            "Direction",
            ["Buy", "Sell"],
            key="trade_direction"
        )
        print(f"Direction selected: {direction}")  # Debug
        
        if ticker:
            print(f"\nProcessing ticker {ticker}")  # Debug
            existing_positions = get_position_summary(ticker)
            print(f"Found existing positions: {existing_positions}")  # Debug
            
            # Handle position selection based on direction
            if direction == "Buy":
                if existing_positions:
                    print(f"Processing existing positions for {ticker}")  # Debug
                    st.write(f"Existing {ticker} Positions:")
                    for pos in existing_positions:
                        print(f"Position details: {pos}")  # Debug
                        st.write(f"Position ID: {pos['position_id']}")
                        st.write(f"Total Shares: {pos['total_shares']}")
                        st.write(f"Average Entry: ${pos['avg_entry_price']:.2f}")
                    
                    add_to_existing = st.checkbox("Add to existing position")
                    print(f"Add to existing selected: {add_to_existing}")  # Debug
                    
                    if add_to_existing:
                        position_id = st.selectbox(
                            "Select Position ID",
                            options=[pos['position_id'] for pos in existing_positions],
                            key="position_select"
                        )
                        print(f"Selected position ID: {position_id}")  # Debug
                    else:
                        position_id = None
                        print("Creating new position")  # Debug
                else:
                    position_id = None
                    print("No existing positions found, creating new position")  # Debug
            else:  # Sell
                if not existing_positions:
                    st.error("No existing positions found for this ticker")
                    st.stop()
                position_id = st.selectbox(
                    "Select Position to Exit",
                    options=[pos['position_id'] for pos in existing_positions],
                    key="position_select"
                )

            # Get entry date/time
            entry_date = st.date_input("Entry Date")
            entry_time_str = st.text_input("Entry Time (HH:MM)", "09:30")
            try:
                entry_time = datetime.strptime(entry_time_str, "%H:%M").time()
                entry_datetime = datetime.combine(entry_date, entry_time)
                entry_datetime = pytz.timezone('US/Pacific').localize(entry_datetime)
            except ValueError:
                st.error("Please enter time in HH:MM format (e.g. 09:30)")
                st.stop()

            # Generate position_id if needed
            if position_id is None:
                position_id = generate_position_id(ticker, entry_datetime)
        
            col1, col2 = st.columns(2)
            with col1:
                shares = st.number_input("Number of Shares", min_value=1, step=1)
                if direction == "Buy":
                    entry_price = st.number_input("Entry Price", min_value=0.01, step=0.01)
                else:
                    entry_price = st.number_input("Exit Price", min_value=0.01, step=0.01)
            
            with col2:
                if direction == "Buy":
                    target_price = st.number_input("Target Price", min_value=0.01, step=0.01)
                    stop_loss = st.number_input("Stop Loss", min_value=0.01, step=0.01)
                    strategy = st.text_input("Strategy")
                else:
                    exit_reason = st.text_area("Exit Reason", key="exit_reason")
                
                order_type = st.selectbox("Order Type", ["Market", "Limit"], key="add_trade_order_type")
            
            
            if direction == "Buy":
                followed_rules = st.checkbox("Followed Trading Rules")
                entry_reason = st.text_area("Entry Reason", key="add_trade_reason")
            
            notes = st.text_area("Notes", key="add_trade_notes")
            
            if st.button("Add Trade"):
                try:
                    print("\n--- Processing Add Trade ---")  # Debug
                    entry_datetime = datetime.combine(entry_date, entry_time)
                    entry_datetime = pytz.timezone('US/Pacific').localize(entry_datetime)
                    print(f"Entry datetime: {entry_datetime}")  # Debug
                    
                    trade = TradeEntry(
                        ticker=ticker,
                        entry_date=entry_datetime,
                        shares=shares,
                        entry_price=entry_price,
                        target_price=target_price if direction == "Buy" else None,
                        stop_loss=stop_loss if direction == "Buy" else None,
                        strategy=strategy if direction == "Buy" else None,
                        order_type=order_type,
                        position_id=position_id,
                        followed_rules=followed_rules if direction == "Buy" else None,
                        entry_reason=entry_reason if direction == "Buy" else None,
                        exit_reason=exit_reason if direction == "Sell" else None,
                        notes=notes,
                        direction=direction.lower()
                    )
                    
                    print(f"Trade object created: {trade.__dict__}")  # Debug
                    
                    add_trade(trade)
                    print("Trade added successfully")  # Debug
                    st.success("Trade added successfully!")
                    st.query_params(rerun=True)
                except Exception as e:
                    print(f"Error adding trade: {str(e)}")  # Debug
                    st.error(f"Error adding trade: {str(e)}")
    
    with tab3:
        st.subheader("Update Trade")
        print("\n--- Update Trade Tab ---")  # Debug
        open_trades = get_open_trades()
        print(f"Retrieved {len(open_trades) if open_trades else 0} trades for update tab")  # Debug
        print(f"Update tab trades  {open_trades}")  # Debug
        
        if open_trades:
            trade_id = st.selectbox(
                "Select Trade to Update",
                options=[t['id'] for t in open_trades],
                format_func=lambda x: f"{next(t['ticker'] for t in open_trades if t['id'] == x)} - {x}",
                key="trade_select"
            )
            
            trade = next(t for t in open_trades if t['id'] == trade_id)
            print(f"\nSelected trade for update: {trade}")  # Debug
            
            col1, col2 = st.columns(2)
            with col1:
                new_shares = st.number_input("Shares", value=trade['shares'] if trade['shares'] is not None else 0)
                new_entry = st.number_input("Entry Price", value=float(trade['entry_price']) if trade['entry_price'] is not None else 0.0)
                new_target = st.number_input("Target Price", value=float(trade['target_price']) if trade['target_price'] is not None else 0.0)
            
            with col2:
                new_stop = st.number_input("Stop Loss", value=float(trade['stop_loss']) if trade['stop_loss'] is not None else 0.0)
                new_strategy = st.text_input("Strategy", value=trade['strategy'])
                new_order_type = st.selectbox("Order Type", ["Market", "Limit"], 
                                            index=0 if trade['order_type'] == "Market" else 1,
                                            key="update_trade_order_type")
            
            # Add date and time fields
            entry_date = st.date_input(
                "Entry Date", 
                value=trade['entry_date'].date(),
                key="update_entry_date"
            )
            entry_time_str = st.text_input(
                "Entry Time (HH:MM)", 
                value=trade['entry_date'].strftime("%H:%M"),
                key="update_entry_time"
            )
            try:
                entry_time = datetime.strptime(entry_time_str, "%H:%M").time()
            except ValueError:
                st.error("Please enter time in HH:MM format (e.g. 09:30)")
                st.stop()
            
            new_notes = st.text_area("Notes", 
                                    value=trade['notes'] if trade['notes'] else "",
                                    key="update_trade_notes")
            
            if st.button("Update Trade"):
                try:
                    # Combine date and time into datetime
                    entry_datetime = datetime.combine(entry_date, entry_time)
                    entry_datetime = pytz.timezone('US/Pacific').localize(entry_datetime)
                    
                    updates = {
                        'entry_date': entry_datetime,
                        'shares': new_shares,
                        'entry_price': new_entry,
                        'target_price': new_target,
                        'stop_loss': new_stop,
                        'strategy': new_strategy,
                        'order_type': new_order_type,
                        'notes': new_notes
                    }
                    
                    print("\n--- Processing Trade Update ---")  # Debug
                    print(f"Update data being sent: {updates}")  # Debug
                    
                    update_trade(trade_id, updates)
                    print("Trade update completed")  # Debug
                    
                    st.success("Trade updated successfully!")
                    print("Clearing session state and triggering rerun")  # Debug
                    
                    # Clear session state
                    for key in st.session_state.keys():
                        del st.session_state[key]
                    
                    st.rerun()
                except Exception as e:
                    st.error(f"Error updating trade: {str(e)}")
        else:
            st.info("No open trades to update")
    
    with tab4:
        st.subheader("Trade History")
        history = get_trade_history()
        
        if history:
            # Group trades by position_id
            positions = {}
            for trade in history:
                if trade['position_id'] not in positions:
                    positions[trade['position_id']] = []
                positions[trade['position_id']].append(trade)

            # Add P/L chart
            fig = plot_trade_history(history)
            if fig:
                st.plotly_chart(fig, use_container_width=True)
            
            # Display trades grouped by position
            for position_id, trades in positions.items():
                # Sort trades chronologically
                trades.sort(key=lambda x: x['entry_date'])
                first_trade = trades[0]
                
                # Calculate position performance
                performance = calculate_position_performance(trades)
                
                with st.expander(f"{first_trade['ticker']} - Position {position_id}"):
                    # Show position summary
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Total Shares Bought", f"{int(performance['total_bought']):,}")
                        st.metric("Avg Entry", f"${performance['avg_entry']:.2f}")
                    with col2:
                        st.metric("Total Shares Sold", f"{int(performance['total_sold']):,}")
                        st.metric("Avg Exit", f"${performance['avg_exit']:.2f}")
                    with col3:
                        st.metric("Remaining Shares", f"{int(performance['remaining_shares']):,}")
                        st.metric("Realized P/L", 
                                f"${performance['realized_pl']:.2f}",
                                f"{(performance['realized_pl'] / (performance['total_sold'] * performance['avg_entry']) * 100):.2f}%" if performance['total_sold'] > 0 else None)
                    
                    st.markdown("---")
                    
                    # Show buy trades
                    st.subheader("Buy Orders")
                    for trade in trades:
                        # Infer direction if not set
                        if 'direction' not in trade or not trade['direction']:
                            is_buy = trade.get('exit_price') is None
                        else:
                            is_buy = trade.get('direction', '').lower() == 'buy'
                            
                        if is_buy:
                            col1, col2, col3 = st.columns(3)
                            with col1:
                                st.text(f"Date: {format_datetime(trade['entry_date'])}")
                                st.text(f"Shares: {trade['shares']}")
                            with col2:
                                st.text(f"Price: ${float(trade['entry_price']):.2f}")
                                st.text(f"Order: {trade['order_type']}")
                            with col3:
                                if trade.get('target_price'):
                                    st.text(f"Target: ${float(trade['target_price']):.2f}")
                                if trade.get('stop_loss'):
                                    st.text(f"Stop: ${float(trade['stop_loss']):.2f}")
                            if trade.get('entry_reason'):
                                st.text(f"Reason: {trade['entry_reason']}")
                            st.markdown("---")
                    
                    # Show sell trades
                    st.subheader("Sell Orders")
                    for trade in trades:
                        # Check for sell orders
                        is_sell = (trade.get('direction', '').lower() == 'sell' or 
                                 trade.get('exit_price') is not None)
                            
                        if is_sell:
                            col1, col2 = st.columns(2)
                            with col1:
                                st.text(f"Date: {format_datetime(trade.get('exit_date') or trade.get('entry_date'))}")
                                st.text(f"Shares: {trade['shares']}")
                            with col2:
                                # Use exit_price for sells, fallback to entry_price
                                price = float(trade.get('exit_price') or trade.get('entry_price'))
                                st.text(f"Price: ${price:.2f}")
                                st.text(f"Order: {trade['order_type']}")
                            if trade.get('exit_reason'):
                                st.text(f"Reason: {trade['exit_reason']}")
                            st.markdown("---")
        else:
            st.info("No trade history found")