import os
import json
import random
import math
import io
from datetime import datetime, timedelta
from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename
from werkzeug.exceptions import RequestEntityTooLarge

app = Flask(__name__)
app.secret_key = os.urandom(24)

# Robustness: Max content length for uploads
app.config['MAX_CONTENT_LENGTH'] = 5 * 1024 * 1024  # 5MB limit
ALLOWED_EXTENSIONS = {'json'}

def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

class GridStrategy:
    def __init__(self, lower_price, upper_price, grid_num, investment, current_price):
        self.lower_price = float(lower_price)
        self.upper_price = float(upper_price)
        self.grid_num = int(grid_num)
        self.investment = float(investment)
        self.initial_price = float(current_price)
        
        # Calculate grid lines
        self.grids = []
        step = (self.upper_price - self.lower_price) / self.grid_num
        for i in range(self.grid_num + 1):
            self.grids.append(self.lower_price + i * step)
            
        self.grid_step = step
        self.cash = self.investment
        self.holdings = 0.0
        self.trades = []
        self.total_arbitrage_profit = 0.0
        
        # Initial Position Building
        # Ideally, if price is in the middle, we should hold some assets to sell as price goes up
        # Simple Logic: Buy assets worth 50% of investment if price is within range
        # Better Logic: Calculate required holdings based on how many grids are ABOVE current price (to sell)
        
        # Let's use a "Geometric" approach or simple "Equal Difference" approach
        # For this demo, we assume we start with 50/50 split if within range
        if self.lower_price < self.initial_price < self.upper_price:
            buy_amount = self.investment / 2
            self.cash -= buy_amount
            self.holdings = buy_amount / self.initial_price
            self.trades.append({
                'type': 'INIT_BUY',
                'price': self.initial_price,
                'amount': self.holdings,
                'time': 'Start'
            })
            
        # Track last grid index crossed
        self.last_grid_index = self._get_grid_index(self.initial_price)

    def _get_grid_index(self, price):
        if price <= self.lower_price:
            return -1
        if price >= self.upper_price:
            return self.grid_num + 1
        return int((price - self.lower_price) / self.grid_step)

    def process_price(self, price, timestamp):
        current_grid_index = self._get_grid_index(price)
        
        # Price moved across grid lines
        if current_grid_index != self.last_grid_index:
            # Check direction and multiple grid crossings
            diff = current_grid_index - self.last_grid_index
            
            # Simple handling: Just process the boundary crossing closest to current price
            # In a real engine, we'd handle gaps. Here we assume granular data or just take the new state.
            
            # Logic:
            # If index increases (price goes up): We crossed a line upwards -> SELL
            # If index decreases (price goes down): We crossed a line downwards -> BUY
            
            # Amount per grid: Total Investment / Grid Num (Simplified)
            amount_per_grid_usdt = self.investment / self.grid_num
            
            action = None
            trade_price = price 
            trade_amount = 0
            
            if diff > 0: 
                # Price Up -> Sell
                # Only sell if we have holdings and we are within valid grid range
                if self.holdings > 0 and 0 <= self.last_grid_index < self.grid_num:
                    # Sell one grid's worth
                    amount_to_sell = amount_per_grid_usdt / price
                    if self.holdings >= amount_to_sell:
                        self.holdings -= amount_to_sell
                        self.cash += amount_to_sell * price
                        trade_amount = amount_to_sell
                        action = 'SELL'
                        
                        # Calculate profit per grid (approximate)
                        # Profit = Buy Price (lower grid) vs Sell Price (this grid)
                        # Roughly = grid_step * amount
                        self.total_arbitrage_profit += self.grid_step * amount_to_sell

            elif diff < 0:
                # Price Down -> Buy
                # Only buy if we have cash and within range
                if self.cash > amount_per_grid_usdt and 0 <= current_grid_index < self.grid_num:
                    amount_to_buy = amount_per_grid_usdt / price
                    self.cash -= amount_to_buy * price
                    self.holdings += amount_to_buy
                    trade_amount = amount_to_buy
                    action = 'BUY'

            if action:
                self.trades.append({
                    'type': action,
                    'price': price,
                    'amount': trade_amount,
                    'time': timestamp,
                    'profit': self.total_arbitrage_profit
                })
            
            self.last_grid_index = current_grid_index

        # Calculate current status
        total_assets_value = self.cash + (self.holdings * price)
        unrealized_pnl = total_assets_value - self.investment
        
        return {
            'price': price,
            'total_value': total_assets_value,
            'cash': self.cash,
            'holdings_value': self.holdings * price,
            'arbitrage_profit': self.total_arbitrage_profit,
            'unrealized_pnl': unrealized_pnl,
            'timestamp': timestamp
        }

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/generate_data', methods=['POST'])
def generate_data():
    try:
        data = request.json
        days = int(data.get('days', 30))
        start_price = float(data.get('start_price', 1000))
        volatility = float(data.get('volatility', 0.02)) # Daily volatility
        trend = float(data.get('trend', 0.000)) # Daily trend
        
        prices = []
        current_price = start_price
        start_date = datetime.now() - timedelta(days=days)
        
        # Generate hourly data
        hours = days * 24
        
        for i in range(hours):
            # Geometric Brownian Motion step
            change = random.normalvariate(trend/24, volatility/math.sqrt(24))
            current_price = current_price * (1 + change)
            timestamp = (start_date + timedelta(hours=i)).strftime('%Y-%m-%d %H:%M')
            prices.append({'time': timestamp, 'price': round(current_price, 2)})
            
        return jsonify({'prices': prices})
    except Exception as e:
        app.logger.error(f"Generate data error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/simulate', methods=['POST'])
def simulate():
    try:
        data = request.json
        prices = data.get('prices', [])
        params = data.get('params', {})
        
        if not prices or not params:
            return jsonify({'error': 'Missing data'}), 400
            
        strategy = GridStrategy(
            lower_price=params.get('lower_price'),
            upper_price=params.get('upper_price'),
            grid_num=params.get('grid_num'),
            investment=params.get('investment'),
            current_price=prices[0]['price']
        )
        
        results = []
        trades = []
        
        for p in prices:
            step_result = strategy.process_price(p['price'], p['time'])
            results.append(step_result)
            
        return jsonify({
            'results': results,
            'trades': strategy.trades,
            'summary': {
                'final_value': results[-1]['total_value'],
                'total_profit': results[-1]['total_value'] - params.get('investment'),
                'arbitrage_profit': strategy.total_arbitrage_profit,
                'grid_yield': (strategy.total_arbitrage_profit / params.get('investment')) * 100
            }
        })
    except Exception as e:
        app.logger.error(f"Simulation error: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/api/upload_config', methods=['POST'])
def upload_config():
    # Robust file upload handling
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
        
    if file and allowed_file(file.filename):
        try:
            # Read file content
            content = file.read()
            
            # Check for binary content (null bytes)
            if b'\0' in content:
                return jsonify({'error': 'Binary file detected'}), 400
                
            # Parse JSON
            try:
                config = json.loads(content.decode('utf-8'))
                return jsonify({'message': 'Config uploaded successfully', 'config': config})
            except UnicodeDecodeError:
                return jsonify({'error': 'File encoding not supported (must be UTF-8)'}), 400
            except json.JSONDecodeError:
                return jsonify({'error': 'Invalid JSON format'}), 400
                
        except Exception as e:
            return jsonify({'error': f"Upload failed: {str(e)}"}), 500
    else:
        return jsonify({'error': 'Invalid file type (only .json allowed)'}), 400

@app.errorhandler(404)
def page_not_found(e):
    # Return a JSON response for API calls, or render a template
    if request.path.startswith('/api/'):
        return jsonify({'error': 'Not found'}), 404
    return render_template('index.html'), 200 # Fallback to SPA entry or create a 404 page

@app.errorhandler(500)
def internal_server_error(e):
    return jsonify({'error': 'Internal Server Error'}), 500

@app.errorhandler(RequestEntityTooLarge)
def handle_file_too_large(e):
    return jsonify({'error': 'File too large (max 5MB)'}), 413

@app.route('/health')
def health():
    return jsonify({"status": "healthy"})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)