"""
Persistent wallet storage with HuggingFace integration
"""
import sqlite3
import time
from typing import Dict, Any, List, Optional
from datetime import datetime
from datasets import load_dataset, Dataset
from huggingface_hub import HfApi, upload_file
import json
import os

class WalletStorage:
    def __init__(self, db_path: str = "wallet_storage.db", 
                 hf_token: Optional[str] = None,
                 dataset_name: str = "my-mining-wallets"):
        self.db_path = db_path
        self.hf_token = hf_token
        self.dataset_name = dataset_name
        self.hf_api = HfApi(token=hf_token)
        self.init_database()

    def init_database(self):
        """Initialize SQLite database for wallets"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()

        # Create wallets table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS wallets (
            address TEXT PRIMARY KEY,
            balance REAL,
            total_mined REAL,
            last_updated TIMESTAMP,
            metadata TEXT
        )
        ''')

        # Create transactions table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS transactions (
            tx_hash TEXT PRIMARY KEY,
            wallet_address TEXT,
            amount REAL,
            timestamp TIMESTAMP,
            block_hash TEXT,
            FOREIGN KEY (wallet_address) REFERENCES wallets(address)
        )
        ''')

        conn.commit()
        conn.close()

    def add_wallet(self, wallet_data: Dict[str, Any]) -> bool:
        """Add or update wallet"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
            INSERT OR REPLACE INTO wallets 
            (address, balance, total_mined, last_updated, metadata)
            VALUES (?, ?, ?, ?, ?)
            ''', (
                wallet_data['address'],
                wallet_data.get('balance', 0.0),
                wallet_data.get('total_mined', 0.0),
                datetime.now(),
                json.dumps(wallet_data.get('metadata', {}))
            ))
            
            conn.commit()
            conn.close()
            
            # Sync with HuggingFace
            self.sync_to_huggingface()
            return True
        except Exception as e:
            print(f"Error adding wallet: {str(e)}")
            return False

    def get_wallet(self, address: str) -> Optional[Dict[str, Any]]:
        """Get wallet details"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
        SELECT * FROM wallets WHERE address = ?
        ''', (address,))
        
        result = cursor.fetchone()
        conn.close()
        
        if result:
            return {
                'address': result[0],
                'balance': result[1],
                'total_mined': result[2],
                'last_updated': result[3],
                'metadata': json.loads(result[4])
            }
        return None

    def update_balance(self, address: str, amount: float) -> bool:
        """Update wallet balance"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
            UPDATE wallets 
            SET balance = balance + ?,
                total_mined = total_mined + ?,
                last_updated = ?
            WHERE address = ?
            ''', (amount, amount, datetime.now(), address))
            
            conn.commit()
            conn.close()
            
            # Sync changes to HuggingFace
            self.sync_to_huggingface()
            return True
        except Exception as e:
            print(f"Error updating balance: {str(e)}")
            return False

    def add_transaction(self, tx_data: Dict[str, Any]) -> bool:
        """Record a new transaction"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
            INSERT INTO transactions 
            (tx_hash, wallet_address, amount, timestamp, block_hash)
            VALUES (?, ?, ?, ?, ?)
            ''', (
                tx_data['tx_hash'],
                tx_data['wallet_address'],
                tx_data['amount'],
                datetime.now(),
                tx_data.get('block_hash', '')
            ))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error adding transaction: {str(e)}")
            return False

    def sync_to_huggingface(self):
        """Sync wallet data to HuggingFace dataset"""
        if not self.hf_token:
            return

        try:
            # Get all wallet data
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('SELECT * FROM wallets')
            wallets = cursor.fetchall()
            
            cursor.execute('SELECT * FROM transactions')
            transactions = cursor.fetchall()
            
            conn.close()

            # Convert to dataset format
            wallet_data = {
                'wallets': [{
                    'address': w[0],
                    'balance': w[1],
                    'total_mined': w[2],
                    'last_updated': w[3],
                    'metadata': w[4]
                } for w in wallets],
                'transactions': [{
                    'tx_hash': t[0],
                    'wallet_address': t[1],
                    'amount': t[2],
                    'timestamp': t[3],
                    'block_hash': t[4]
                } for t in transactions]
            }

            # Save to temporary file
            temp_file = 'wallet_data.json'
            with open(temp_file, 'w') as f:
                json.dump(wallet_data, f)

            # Upload to HuggingFace
            self.hf_api.upload_file(
                path_or_fileobj=temp_file,
                path_in_repo='wallet_data.json',
                repo_id=self.dataset_name,
                repo_type='dataset'
            )

            # Clean up
            os.remove(temp_file)

        except Exception as e:
            print(f"Error syncing to HuggingFace: {str(e)}")

    def load_from_huggingface(self):
        """Load wallet data from HuggingFace dataset"""
        if not self.hf_token:
            return

        try:
            # Load dataset from HuggingFace
            dataset = load_dataset(self.dataset_name)
            
            if 'wallet_data.json' in dataset:
                wallet_data = dataset['wallet_data.json']
                
                # Update local database
                conn = sqlite3.connect(self.db_path)
                cursor = conn.cursor()
                
                # Update wallets
                for wallet in wallet_data['wallets']:
                    cursor.execute('''
                    INSERT OR REPLACE INTO wallets 
                    (address, balance, total_mined, last_updated, metadata)
                    VALUES (?, ?, ?, ?, ?)
                    ''', (
                        wallet['address'],
                        wallet['balance'],
                        wallet['total_mined'],
                        wallet['last_updated'],
                        wallet['metadata']
                    ))
                
                # Update transactions
                for tx in wallet_data['transactions']:
                    cursor.execute('''
                    INSERT OR IGNORE INTO transactions 
                    (tx_hash, wallet_address, amount, timestamp, block_hash)
                    VALUES (?, ?, ?, ?, ?)
                    ''', (
                        tx['tx_hash'],
                        tx['wallet_address'],
                        tx['amount'],
                        tx['timestamp'],
                        tx['block_hash']
                    ))
                
                conn.commit()
                conn.close()
                
        except Exception as e:
            print(f"Error loading from HuggingFace: {str(e)}")
