"""
Wallet Database Manager
Handles wallet states, balances, transactions, and transfers
"""
import duckdb
import time
from typing import List, Dict, Optional
from decimal import Decimal
import json
from datetime import datetime

class WalletDB:
    def __init__(self, db_path: str = "data/wallet/wallet_storage.db"):
        """Initialize wallet database with local storage
        Args:
            db_path: Path to the local database file. Will be created if it doesn't exist.
        """
        # Ensure the directory exists
        import os
        os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)
        
        self.db_path = db_path
        self.conn = self._init_db_connection()
        self.setup_tables()

    def _init_db_connection(self) -> duckdb.DuckDBPyConnection:
        """Initialize database connection with local storage configuration"""
        try:
            # Create connection with local file
            conn = duckdb.connect(self.db_path)
            
            # Enable WAL mode for better concurrent access
            conn.execute("PRAGMA journal_mode=WAL;")
            
            return conn
        except Exception as e:
            raise Exception(f"Failed to initialize database connection: {str(e)}")

    def setup_tables(self):
        """Create necessary database tables"""
        # Wallets table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS wallets (
                address VARCHAR PRIMARY KEY,
                public_key VARCHAR NOT NULL,
                encrypted_private_key VARCHAR,
                wif_private_key VARCHAR,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_accessed TIMESTAMP,
                is_active BOOLEAN DEFAULT true
            )
        """)

        # Balances table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS balances (
                address VARCHAR,
                currency VARCHAR,
                amount DECIMAL(24,8) DEFAULT 0,
                last_updated TIMESTAMP,
                PRIMARY KEY (address, currency),
                FOREIGN KEY (address) REFERENCES wallets(address)
            )
        """)

        # Transactions table
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS transactions (
                tx_hash VARCHAR PRIMARY KEY,
                from_address VARCHAR,
                to_address VARCHAR,
                amount DECIMAL(24,8),
                currency VARCHAR,
                fee DECIMAL(24,8),
                status VARCHAR,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                block_height INTEGER,
                confirmations INTEGER DEFAULT 0,
                raw_tx TEXT,
                FOREIGN KEY (from_address) REFERENCES wallets(address),
                FOREIGN KEY (to_address) REFERENCES wallets(address)
            )
        """)

        # Transaction Status History
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS tx_status_history (
                tx_hash VARCHAR,
                status VARCHAR,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                details TEXT,
                FOREIGN KEY (tx_hash) REFERENCES transactions(tx_hash)
            )
        """)

        # Address Book
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS address_book (
                owner_address VARCHAR,
                contact_address VARCHAR,
                label VARCHAR,
                notes TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (owner_address) REFERENCES wallets(address)
            )
        """)

    def save_wallet(self, wallet_info: Dict) -> bool:
        """Save wallet information to database"""
        try:
            self.conn.execute("""
                INSERT INTO wallets (address, public_key, encrypted_private_key, wif_private_key)
                VALUES (?, ?, ?, ?)
            """, [
                wallet_info['address'],
                wallet_info['public_key'],
                wallet_info.get('encrypted_private_key'),
                wallet_info['wif_private_key']
            ])

            # Initialize balance for BTC
            self.conn.execute("""
                INSERT INTO balances (address, currency, amount, last_updated)
                VALUES (?, 'BTC', 0, CURRENT_TIMESTAMP)
            """, [wallet_info['address']])

            return True
        except Exception as e:
            print(f"Error saving wallet: {str(e)}")
            return False

    def get_wallet(self, address: str) -> Optional[Dict]:
        """Retrieve wallet information"""
        result = self.conn.execute("""
            SELECT w.*, b.amount as btc_balance
            FROM wallets w
            LEFT JOIN balances b ON w.address = b.address AND b.currency = 'BTC'
            WHERE w.address = ?
        """, [address]).fetchone()

        if result:
            self.update_last_accessed(address)
            return {
                'address': result[0],
                'public_key': result[1],
                'encrypted_private_key': result[2],
                'wif_private_key': result[3],
                'created_at': result[4],
                'last_accessed': result[5],
                'is_active': result[6],
                'btc_balance': float(result[7]) if result[7] else 0
            }
        return None

    def update_balance(self, address: str, currency: str, amount: Decimal) -> bool:
        """Update wallet balance"""
        try:
            self.conn.execute("""
                INSERT INTO balances (address, currency, amount, last_updated)
                VALUES (?, ?, ?, CURRENT_TIMESTAMP)
                ON CONFLICT (address, currency) DO UPDATE
                SET amount = ?, last_updated = CURRENT_TIMESTAMP
            """, [address, currency, amount, amount])
            return True
        except Exception as e:
            print(f"Error updating balance: {str(e)}")
            return False

    def record_transaction(self, tx_data: Dict) -> bool:
        """Record a new transaction"""
        try:
            # Insert transaction
            self.conn.execute("""
                INSERT INTO transactions (
                    tx_hash, from_address, to_address, amount, 
                    currency, fee, status, raw_tx
                )
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, [
                tx_data['tx_hash'],
                tx_data['from_address'],
                tx_data['to_address'],
                tx_data['amount'],
                tx_data.get('currency', 'BTC'),
                tx_data.get('fee', 0),
                'pending',
                json.dumps(tx_data.get('raw_tx', {}))
            ])

            # Record initial status
            self.conn.execute("""
                INSERT INTO tx_status_history (tx_hash, status, details)
                VALUES (?, 'pending', ?)
            """, [tx_data['tx_hash'], 'Transaction initiated'])

            return True
        except Exception as e:
            print(f"Error recording transaction: {str(e)}")
            return False

    def update_transaction_status(self, tx_hash: str, status: str, details: str = None) -> bool:
        """Update transaction status"""
        try:
            self.conn.execute("""
                UPDATE transactions 
                SET status = ?, 
                    confirmations = CASE 
                        WHEN ? = 'confirmed' THEN confirmations + 1 
                        ELSE confirmations 
                    END
                WHERE tx_hash = ?
            """, [status, status, tx_hash])

            # Record status change
            self.conn.execute("""
                INSERT INTO tx_status_history (tx_hash, status, details)
                VALUES (?, ?, ?)
            """, [tx_hash, status, details or f"Status updated to {status}"])

            return True
        except Exception as e:
            print(f"Error updating transaction status: {str(e)}")
            return False

    def get_transaction_history(self, address: str) -> List[Dict]:
        """Get transaction history for an address"""
        results = self.conn.execute("""
            SELECT t.*, th.details
            FROM transactions t
            LEFT JOIN tx_status_history th ON t.tx_hash = th.tx_hash
            WHERE t.from_address = ? OR t.to_address = ?
            ORDER BY t.timestamp DESC
        """, [address, address]).fetchall()

        return [{
            'tx_hash': row[0],
            'from_address': row[1],
            'to_address': row[2],
            'amount': float(row[3]),
            'currency': row[4],
            'fee': float(row[5]),
            'status': row[6],
            'timestamp': row[7],
            'confirmations': row[9],
            'details': row[-1]
        } for row in results]

    def add_address_book_entry(self, owner_address: str, contact_address: str, 
                             label: str, notes: str = None) -> bool:
        """Add an address book entry"""
        try:
            self.conn.execute("""
                INSERT INTO address_book (owner_address, contact_address, label, notes)
                VALUES (?, ?, ?, ?)
            """, [owner_address, contact_address, label, notes])
            return True
        except Exception as e:
            print(f"Error adding address book entry: {str(e)}")
            return False

    def get_address_book(self, owner_address: str) -> List[Dict]:
        """Get address book entries for a wallet"""
        results = self.conn.execute("""
            SELECT contact_address, label, notes, created_at
            FROM address_book
            WHERE owner_address = ?
            ORDER BY label
        """, [owner_address]).fetchall()

        return [{
            'address': row[0],
            'label': row[1],
            'notes': row[2],
            'created_at': row[3]
        } for row in results]

    def update_last_accessed(self, address: str) -> None:
        """Update last accessed timestamp for a wallet"""
        self.conn.execute("""
            UPDATE wallets 
            SET last_accessed = CURRENT_TIMESTAMP 
            WHERE address = ?
        """, [address])

    def close(self):
        """Close database connection"""
        self.conn.close()
