import sqlite3
import pandas as pd
import yfinance as yf

DEMO_MONEY = 100000

# Initialize SQLite DB
def init_db():
    conn = sqlite3.connect("portfolio.db")
    conn.execute('''
        CREATE TABLE IF NOT EXISTS trades (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            symbol TEXT,
            quantity INTEGER,
            price REAL,
            total REAL,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
        )
    ''')
    conn.commit()
    conn.close()

init_db()

# Get current balance
def get_balance():
    conn = sqlite3.connect("portfolio.db")
    df = pd.read_sql("SELECT SUM(total) as spent FROM trades WHERE quantity > 0", conn)
    conn.close()
    spent = df['spent'][0] if df['spent'][0] else 0
    return DEMO_MONEY - spent

# Get how many stocks owned
def get_owned_quantity(symbol):
    conn = sqlite3.connect("portfolio.db")
    df = pd.read_sql(f"SELECT SUM(quantity) as qty FROM trades WHERE symbol = '{symbol}'", conn)
    conn.close()
    return df['qty'][0] if df['qty'][0] else 0

# Buy stock
def buy_stock(symbol, quantity):
    try:
        price = yf.download(symbol, period="1d", interval="1m")['Close'].iloc[-1]
    except:
        return "❌ Could not fetch live price."

    total = quantity * price
    balance = get_balance()
    if total > balance:
        return f"❌ Not enough balance. You have ₹{balance:.2f}"

    conn = sqlite3.connect("portfolio.db")
    conn.execute("INSERT INTO trades (symbol, quantity, price, total) VALUES (?, ?, ?, ?)",
                 (symbol, quantity, price, total))
    conn.commit()
    conn.close()
    return f"✅ Bought {quantity} of {symbol} at ₹{price:.2f}"

# Sell stock
def sell_stock(symbol, quantity):
    owned = get_owned_quantity(symbol)
    if quantity > owned:
        return f"❌ You only own {owned} shares of {symbol}"

    try:
        price = yf.download(symbol, period="1d", interval="1m")['Close'].iloc[-1]
    except:
        return "❌ Could not fetch price."

    total = quantity * price
    conn = sqlite3.connect("portfolio.db")
    conn.execute("INSERT INTO trades (symbol, quantity, price, total) VALUES (?, ?, ?, ?)",
                 (symbol, -quantity, price, -total))
    conn.commit()
    conn.close()
    return f"✅ Sold {quantity} of {symbol} at ₹{price:.2f}"

# Portfolio table
def get_portfolio():
    conn = sqlite3.connect("portfolio.db")
    df = pd.read_sql("SELECT symbol, SUM(quantity) as qty FROM trades GROUP BY symbol", conn)
    conn.close()

    rows = []
    for _, row in df.iterrows():
        if row['qty'] > 0:
            live_price = yf.download(row['symbol'], period="1d", interval="1m")['Close'].iloc[-1]
            total_value = live_price * row['qty']
            rows.append([row['symbol'], row['qty'], round(live_price, 2), round(total_value, 2)])

    return pd.DataFrame(rows, columns=["Symbol", "Quantity", "Live Price", "Value"])

# Net worth
def get_net_worth():
    cash = get_balance()
    port_df = get_portfolio()
    stock_value = port_df["Value"].sum() if not port_df.empty else 0
    return round(cash + stock_value, 2)