import aiosqlite
import json
from datetime import datetime

DATABASE_FILE = "app_data.db"

async def init_db():
    async with aiosqlite.connect(DATABASE_FILE) as db:
        await db.executescript("""
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                email TEXT UNIQUE NOT NULL,
                api_key TEXT UNIQUE,
                plan TEXT DEFAULT 'free'
            );

            CREATE TABLE IF NOT EXISTS wallets (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                address TEXT NOT NULL,
                name TEXT,
                FOREIGN KEY (user_id) REFERENCES users(id)
            );

            CREATE TABLE IF NOT EXISTS transactions (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                wallet_id INTEGER,
                timestamp TEXT,
                type TEXT,
                sell_currency TEXT,
                sell_amount REAL,
                sell_value REAL,
                buy_currency TEXT,
                buy_amount REAL,
                buy_value REAL,
                closed BOOLEAN DEFAULT 0,
                details TEXT,
                FOREIGN KEY (wallet_id) REFERENCES wallets(id)
            );

            CREATE TABLE IF NOT EXISTS holdings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                wallet_id INTEGER,
                currency TEXT,
                amount REAL,
                last_updated TEXT,
                FOREIGN KEY (wallet_id) REFERENCES wallets(id)
            );

            CREATE TABLE IF NOT EXISTS price_alerts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                currency TEXT,
                target_price REAL,
                alert_type TEXT,
                FOREIGN KEY (user_id) REFERENCES users(id)
            );

            CREATE TABLE IF NOT EXISTS followed_accounts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                address TEXT,
                followed_address TEXT,
                name TEXT,
                FOREIGN KEY (address) REFERENCES wallets(address),
                FOREIGN KEY (followed_address) REFERENCES wallets(address),
                FOREIGN KEY (user_id) REFERENCES users(id)
            );
        """)
        await db.commit()

async def store_transaction(wallet_id, transaction_type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, details=None):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        await db.execute("""
            INSERT INTO transactions (wallet_id, timestamp, type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, details)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (wallet_id, datetime.now().isoformat(), transaction_type, sell_currency, sell_amount, sell_value, buy_currency, buy_amount, buy_value, json.dumps(details or {})))
        await db.commit()

async def update_holdings(wallet_id, currency, amount_change):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("SELECT amount FROM holdings WHERE wallet_id = ? AND currency = ?", (wallet_id, currency))
        result = await cursor.fetchone()
        if result:
            new_amount = result[0] + amount_change
            await db.execute("UPDATE holdings SET amount = ?, last_updated = ? WHERE wallet_id = ? AND currency = ?",
                             (new_amount, datetime.now().isoformat(), wallet_id, currency))
        else:
            await db.execute("INSERT INTO holdings (wallet_id, currency, amount, last_updated) VALUES (?, ?, ?, ?)",
                             (wallet_id, currency, amount_change, datetime.now().isoformat()))
        await db.commit()

async def get_wallet_holdings(wallet_id):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("SELECT currency, amount FROM holdings WHERE wallet_id = ?", (wallet_id,))
        return await cursor.fetchall()

async def get_transaction_history(wallet_id, start_date=None, end_date=None, include_closed=False):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        query = "SELECT * FROM transactions WHERE wallet_id = ?"
        params = [wallet_id]
        if not include_closed:
            query += " AND closed = 0"
        if start_date:
            query += " AND timestamp >= ?"
            params.append(start_date)
        if end_date:
            query += " AND timestamp <= ?"
            params.append(end_date)
        query += " ORDER BY timestamp DESC"
        cursor = await db.execute(query, params)
        return await cursor.fetchall()

# New utility functions

async def close_transaction(transaction_id):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        await db.execute("UPDATE transactions SET closed = 1 WHERE id = ?", (transaction_id,))
        await db.commit()

async def get_open_transactions(wallet_id, currency):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("""
            SELECT * FROM transactions 
            WHERE wallet_id = ? AND buy_currency = ? AND closed = 0
            ORDER BY timestamp ASC
        """, (wallet_id, currency))
        return await cursor.fetchall()

async def calculate_current_holdings(wallet_id):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("""
            SELECT 
                buy_currency AS currency,
                SUM(buy_amount) - COALESCE(
                    (SELECT SUM(sell_amount) 
                     FROM transactions t2 
                     WHERE t2.wallet_id = t1.wallet_id 
                     AND t2.sell_currency = t1.buy_currency 
                     AND t2.closed = 0),
                    0
                ) AS amount
            FROM transactions t1
            WHERE wallet_id = ? AND closed = 0
            GROUP BY buy_currency
            HAVING amount > 0
        """, (wallet_id,))
        return await cursor.fetchall()

STABLECOINS = ['USDC', 'USDT', 'SOL']

async def is_transaction_closed(wallet_id, transaction_id):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("""
            SELECT t1.buy_currency, t1.buy_amount,
                   (SELECT SUM(sell_amount) 
                    FROM transactions t2 
                    WHERE t2.wallet_id = t1.wallet_id 
                    AND t2.sell_currency = t1.buy_currency 
                    AND t2.timestamp > t1.timestamp) AS sold_amount
            FROM transactions t1
            WHERE t1.id = ? AND t1.wallet_id = ?
        """, (transaction_id, wallet_id))
        result = await cursor.fetchone()
        
        if result:
            buy_currency, buy_amount, sold_amount = result
            return sold_amount is not None and sold_amount >= buy_amount
        return False

async def close_completed_transactions(wallet_id):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        cursor = await db.execute("""
            SELECT id FROM transactions
            WHERE wallet_id = ? AND closed = 0 AND buy_currency NOT IN (?)
        """, (wallet_id, ','.join(STABLECOINS)))
        transactions = await cursor.fetchall()
        
        for (transaction_id,) in transactions:
            if await is_transaction_closed(wallet_id, transaction_id):
                await close_transaction(transaction_id)

async def get_profit_loss(wallet_id, currency, start_date=None, end_date=None):
    async with aiosqlite.connect(DATABASE_FILE) as db:
        query = """
            SELECT 
                SUM(CASE WHEN sell_currency = ? THEN sell_value ELSE -buy_value END) as profit_loss
            FROM transactions
            WHERE wallet_id = ? AND (sell_currency = ? OR buy_currency = ?)
        """
        params = [currency, wallet_id, currency, currency]
        
        if start_date:
            query += " AND timestamp >= ?"
            params.append(start_date)
        if end_date:
            query += " AND timestamp <= ?"
            params.append(end_date)
        
        cursor = await db.execute(query, params)
        result = await cursor.fetchone()
        return result[0] if result else 0

# Example usage
if __name__ == "__main__":
    import asyncio

    async def main():
        await init_db()
        # Add more test functions here

    asyncio.run(main())