from datetime import datetime
from dataclasses import dataclass
from typing import Optional
from db.db_connection import create_client
from trading.position_calculator import PositionCalculator

@dataclass
class TradeEntry:
    ticker: str
    entry_date: datetime
    shares: int
    entry_price: float
    target_price: float
    stop_loss: float
    strategy: str
    order_type: str  # New field for Market/Limit
    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  # New field to group related orders

    @property
    def expected_profit_loss(self) -> float:
        return (self.target_price - self.entry_price) * self.shares

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

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 Float64,
                stop_loss Float64,
                strategy String,
                order_type 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()
            ) ENGINE = MergeTree()
            ORDER BY (position_id, id, entry_date)
        """
        client.execute(query)

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

def generate_position_id(ticker: str) -> str:
    """Generate a unique position ID for grouping related trades"""
    timestamp = datetime.now().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"""
    with create_client() as client:
        query = f"""
            SELECT 
                position_id,
                sum(shares) as total_shares,
                sum(shares * entry_price) / sum(shares) as avg_entry_price,
                min(entry_date) as first_entry,
                max(entry_date) as last_entry,
                count() as num_orders
            FROM stock_db.trades 
            WHERE ticker = '{ticker}' 
            AND exit_price IS NULL 
            GROUP BY position_id
            ORDER BY first_entry DESC
        """
        result = client.execute(query)
        columns = ['position_id', 'total_shares', 'avg_entry_price', 
                  'first_entry', 'last_entry', 'num_orders']
        return [dict(zip(columns, row)) for row in result]

def get_order_type() -> str:
    """Get order type from user"""
    while True:
        print("\nOrder Type:")
        print("1. Market")
        print("2. Limit")
        choice = input("Select order type (1-2): ")
        if choice == "1":
            return "Market"
        elif choice == "2":
            return "Limit"
        print("Invalid choice. Please try again.")

def add_trade(trade: TradeEntry):
    with create_client() as client:
        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 (
                {generate_id()},
                '{trade.position_id}',
                '{trade.ticker}',
                '{trade.entry_date.strftime('%Y-%m-%d %H:%M:%S')}',
                {trade.shares},
                {trade.entry_price},
                {trade.target_price},
                {trade.stop_loss},
                '{trade.strategy}',
                '{trade.order_type}',
                {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.execute(query)

def update_trade_exit(trade_id: int, exit_price: float, exit_date: datetime,
                     followed_rules: bool, exit_reason: str, notes: Optional[str] = None):
    with create_client() as client:
        query = f"""
            ALTER TABLE stock_db.trades
            UPDATE 
                exit_price = {exit_price},
                exit_date = '{exit_date.strftime('%Y-%m-%d %H:%M:%S')}',
                followed_rules = {1 if followed_rules else 0},
                exit_reason = {f"'{exit_reason}'" if exit_reason else 'NULL'},
                notes = {f"'{notes}'" if notes else 'notes'}
            WHERE id = {trade_id}
        """
        client.execute(query)

def get_open_trades():
    with create_client() as client:
        query = "SELECT * FROM stock_db.trades WHERE exit_price IS NULL ORDER BY entry_date DESC"
        result = client.execute(query)
        columns = ['id', 'ticker', 'entry_date', 'shares', 'entry_price', 'target_price', 
                  'stop_loss', 'strategy', 'followed_rules', 'entry_reason', 'exit_price', 
                  'exit_date', 'exit_reason', 'notes', 'created_at']
        return [dict(zip(columns, row)) for row in result]

def get_trade_history(limit: int = 50):
    with create_client() as client:
        query = f"""
            SELECT * FROM stock_db.trades 
            WHERE exit_price IS NOT NULL 
            ORDER BY exit_date DESC 
            LIMIT {limit}
        """
        result = client.execute(query)
        columns = ['id', 'ticker', 'entry_date', 'shares', 'entry_price', 'target_price', 
                  'stop_loss', 'strategy', 'followed_rules', 'entry_reason', 'exit_price', 
                  'exit_date', 'exit_reason', 'notes', 'created_at']
        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. Return to Main Menu")
        
        choice = input("\nSelect an option (1-5): ")
        
        if choice == "1":
            ticker = input("Enter ticker symbol: ").upper()
            
            # 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 = input("\nAdd to existing position? (y/n): ").lower() == 'y'
                if add_to_existing:
                    position_id = input("Enter Position ID: ")
                else:
                    position_id = generate_position_id(ticker)
            else:
                position_id = generate_position_id(ticker)
            
            shares = int(input("Enter number of shares: "))
            entry_price = float(input("Enter entry price: "))
            order_type = get_order_type()
            
            # 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=datetime.now(),
                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']}")
            
            trade_id = int(input("\nEnter trade ID to update: "))
            exit_price = float(input("Enter exit price: "))
            followed_rules = input("Did you follow your rules? (y/n): ").lower() == 'y'
            exit_reason = input("Enter exit reason: ")
            notes = input("Additional notes (optional): ") or None
            
            update_trade_exit(
                trade_id=trade_id,
                exit_price=exit_price,
                exit_date=datetime.now(),
                followed_rules=followed_rules,
                exit_reason=exit_reason,
                notes=notes
            )
            print("Trade updated successfully!")
            
        elif choice == "3":
            open_trades = get_open_trades()
            if not open_trades:
                print("No open trades found.")
            else:
                print("\nOpen Trades:")
                for trade in open_trades:
                    print(f"\nTicker: {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']}")
                    
        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":
            break