"""
Atomizer Dashboard API

RESTful API for controlling and monitoring optimization runs.
Provides endpoints for:
- Starting/stopping optimizations
- Managing studies (list, resume, delete)
- Real-time monitoring of progress
- Retrieving results and visualizations
"""

from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS
import json
import sys
from pathlib import Path
from typing import Dict, List, Any
import threading
import time
from datetime import datetime

# Add project root to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from optimization_engine.runner import OptimizationRunner

app = Flask(__name__, static_folder='../frontend', static_url_path='')
CORS(app)

# Global state for running optimizations
active_optimizations = {}
optimization_lock = threading.Lock()


@app.route('/')
def index():
    """Serve the dashboard frontend."""
    return send_from_directory(app.static_folder, 'index.html')


@app.route('/api/studies', methods=['GET'])
def list_studies():
    """
    List all available optimization studies.

    Returns:
        JSON array of study metadata
    """
    try:
        # Use a dummy runner to access list_studies
        config_path = project_root / 'examples/bracket/optimization_config_stress_displacement.json'
        runner = OptimizationRunner(
            config_path=config_path,
            model_updater=lambda x: None,
            simulation_runner=lambda: Path('dummy.op2'),
            result_extractors={}
        )

        studies = runner.list_studies()
        return jsonify({
            'success': True,
            'studies': studies
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/studies/<study_name>', methods=['GET'])
def get_study_details(study_name: str):
    """
    Get detailed information about a specific study.

    Args:
        study_name: Name of the study

    Returns:
        JSON with study metadata, history, and summary
    """
    try:
        config_path = project_root / 'examples/bracket/optimization_config_stress_displacement.json'
        runner = OptimizationRunner(
            config_path=config_path,
            model_updater=lambda x: None,
            simulation_runner=lambda: Path('dummy.op2'),
            result_extractors={}
        )

        output_dir = runner.output_dir

        # Load history
        history_path = output_dir / 'history.json'
        history = []
        if history_path.exists():
            with open(history_path, 'r') as f:
                history = json.load(f)

        # Load summary
        summary_path = output_dir / 'optimization_summary.json'
        summary = {}
        if summary_path.exists():
            with open(summary_path, 'r') as f:
                summary = json.load(f)

        # Load metadata
        metadata_path = runner._get_study_metadata_path(study_name)
        metadata = {}
        if metadata_path.exists():
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)

        return jsonify({
            'success': True,
            'study_name': study_name,
            'metadata': metadata,
            'history': history,
            'summary': summary
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/studies/<study_name>/delete', methods=['DELETE'])
def delete_study(study_name: str):
    """
    Delete a study and all its data.

    Args:
        study_name: Name of the study to delete
    """
    try:
        config_path = project_root / 'examples/bracket/optimization_config_stress_displacement.json'
        runner = OptimizationRunner(
            config_path=config_path,
            model_updater=lambda x: None,
            simulation_runner=lambda: Path('dummy.op2'),
            result_extractors={}
        )

        # Delete database and metadata
        db_path = runner._get_study_db_path(study_name)
        metadata_path = runner._get_study_metadata_path(study_name)

        if db_path.exists():
            db_path.unlink()
        if metadata_path.exists():
            metadata_path.unlink()

        return jsonify({
            'success': True,
            'message': f'Study {study_name} deleted successfully'
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/optimization/start', methods=['POST'])
def start_optimization():
    """
    Start a new optimization run or resume an existing one.

    Request body:
        {
            "study_name": "my_study",
            "n_trials": 50,
            "resume": false,
            "config_path": "path/to/config.json"
        }
    """
    try:
        data = request.get_json()
        study_name = data.get('study_name', f'study_{datetime.now().strftime("%Y%m%d_%H%M%S")}')
        n_trials = data.get('n_trials', 50)
        resume = data.get('resume', False)
        config_path = data.get('config_path', 'examples/bracket/optimization_config_stress_displacement.json')

        with optimization_lock:
            if study_name in active_optimizations:
                return jsonify({
                    'success': False,
                    'error': f'Study {study_name} is already running'
                }), 400

            # Mark as active
            active_optimizations[study_name] = {
                'status': 'starting',
                'start_time': datetime.now().isoformat(),
                'n_trials': n_trials,
                'current_trial': 0
            }

        # Start optimization in background thread
        def run_optimization():
            try:
                # Import necessary functions
                from examples.test_journal_optimization import (
                    bracket_model_updater,
                    bracket_simulation_runner,
                    stress_extractor,
                    displacement_extractor
                )

                runner = OptimizationRunner(
                    config_path=project_root / config_path,
                    model_updater=bracket_model_updater,
                    simulation_runner=bracket_simulation_runner,
                    result_extractors={
                        'stress_extractor': stress_extractor,
                        'displacement_extractor': displacement_extractor
                    }
                )

                with optimization_lock:
                    active_optimizations[study_name]['status'] = 'running'

                study = runner.run(
                    study_name=study_name,
                    n_trials=n_trials,
                    resume=resume
                )

                with optimization_lock:
                    active_optimizations[study_name]['status'] = 'completed'
                    active_optimizations[study_name]['end_time'] = datetime.now().isoformat()
                    active_optimizations[study_name]['best_value'] = study.best_value
                    active_optimizations[study_name]['best_params'] = study.best_params

            except Exception as e:
                with optimization_lock:
                    active_optimizations[study_name]['status'] = 'failed'
                    active_optimizations[study_name]['error'] = str(e)

        thread = threading.Thread(target=run_optimization, daemon=True)
        thread.start()

        return jsonify({
            'success': True,
            'message': f'Optimization {study_name} started',
            'study_name': study_name
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/optimization/status', methods=['GET'])
def get_optimization_status():
    """
    Get status of all active optimizations.

    Returns:
        JSON with status of all running/recent optimizations
    """
    with optimization_lock:
        return jsonify({
            'success': True,
            'active_optimizations': active_optimizations
        })


@app.route('/api/optimization/<study_name>/status', methods=['GET'])
def get_study_status(study_name: str):
    """
    Get real-time status of a specific optimization.

    Args:
        study_name: Name of the study
    """
    with optimization_lock:
        if study_name not in active_optimizations:
            # Try to get from history
            try:
                config_path = project_root / 'examples/bracket/optimization_config_stress_displacement.json'
                runner = OptimizationRunner(
                    config_path=config_path,
                    model_updater=lambda x: None,
                    simulation_runner=lambda: Path('dummy.op2'),
                    result_extractors={}
                )

                history_path = runner.output_dir / 'history.json'
                if history_path.exists():
                    with open(history_path, 'r') as f:
                        history = json.load(f)

                    return jsonify({
                        'success': True,
                        'status': 'completed',
                        'n_trials': len(history)
                    })
            except:
                pass

            return jsonify({
                'success': False,
                'error': 'Study not found'
            }), 404

        return jsonify({
            'success': True,
            **active_optimizations[study_name]
        })


@app.route('/api/config/load', methods=['GET'])
def load_config():
    """
    Load optimization configuration from file.

    Query params:
        path: Path to config file (relative to project root)
    """
    try:
        config_path = request.args.get('path', 'examples/bracket/optimization_config_stress_displacement.json')
        full_path = project_root / config_path

        with open(full_path, 'r') as f:
            config = json.load(f)

        return jsonify({
            'success': True,
            'config': config,
            'path': config_path
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/config/save', methods=['POST'])
def save_config():
    """
    Save optimization configuration to file.

    Request body:
        {
            "path": "path/to/config.json",
            "config": {...}
        }
    """
    try:
        data = request.get_json()
        config_path = data.get('path')
        config = data.get('config')

        if not config_path or not config:
            return jsonify({
                'success': False,
                'error': 'Missing path or config'
            }), 400

        full_path = project_root / config_path
        full_path.parent.mkdir(parents=True, exist_ok=True)

        with open(full_path, 'w') as f:
            json.dump(config, f, indent=2)

        return jsonify({
            'success': True,
            'message': f'Configuration saved to {config_path}'
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/results/visualization/<study_name>', methods=['GET'])
def get_visualization_data(study_name: str):
    """
    Get data formatted for visualization (charts).

    Args:
        study_name: Name of the study
    """
    try:
        config_path = project_root / 'examples/bracket/optimization_config_stress_displacement.json'
        runner = OptimizationRunner(
            config_path=config_path,
            model_updater=lambda x: None,
            simulation_runner=lambda: Path('dummy.op2'),
            result_extractors={}
        )

        history_path = runner.output_dir / 'history.json'
        if not history_path.exists():
            return jsonify({
                'success': False,
                'error': 'No history found for this study'
            }), 404

        with open(history_path, 'r') as f:
            history = json.load(f)

        # Format data for charts
        trials = [entry['trial_number'] for entry in history]
        objectives = {}
        design_vars = {}
        constraints = {}

        for entry in history:
            for obj_name, obj_value in entry['objectives'].items():
                if obj_name not in objectives:
                    objectives[obj_name] = []
                objectives[obj_name].append(obj_value)

            for dv_name, dv_value in entry['design_variables'].items():
                if dv_name not in design_vars:
                    design_vars[dv_name] = []
                design_vars[dv_name].append(dv_value)

            for const_name, const_value in entry['constraints'].items():
                if const_name not in constraints:
                    constraints[const_name] = []
                constraints[const_name].append(const_value)

        # Calculate running best
        total_objectives = [entry['total_objective'] for entry in history]
        running_best = []
        current_best = float('inf')
        for val in total_objectives:
            current_best = min(current_best, val)
            running_best.append(current_best)

        return jsonify({
            'success': True,
            'trials': trials,
            'objectives': objectives,
            'design_variables': design_vars,
            'constraints': constraints,
            'total_objectives': total_objectives,
            'running_best': running_best
        })
    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


# ====================
# Study Management API
# ====================

@app.route('/api/study/create', methods=['POST'])
def create_study():
    """
    Create a new study with folder structure.

    Request body:
        {
            "study_name": "my_new_study",
            "description": "Optional description"
        }
    """
    try:
        data = request.get_json()
        study_name = data.get('study_name')
        description = data.get('description', '')

        if not study_name:
            return jsonify({
                'success': False,
                'error': 'study_name is required'
            }), 400

        # Create study folder structure
        study_dir = project_root / 'optimization_results' / study_name

        if study_dir.exists():
            return jsonify({
                'success': False,
                'error': f'Study {study_name} already exists'
            }), 400

        # Create directories
        study_dir.mkdir(parents=True, exist_ok=True)
        (study_dir / 'sim').mkdir(exist_ok=True)
        (study_dir / 'results').mkdir(exist_ok=True)

        # Create initial metadata
        metadata = {
            'study_name': study_name,
            'description': description,
            'created_at': datetime.now().isoformat(),
            'status': 'created',
            'has_sim_files': False,
            'is_configured': False
        }

        metadata_path = study_dir / 'metadata.json'
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)

        return jsonify({
            'success': True,
            'message': f'Study {study_name} created successfully',
            'study_path': str(study_dir),
            'sim_folder': str(study_dir / 'sim')
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/study/<study_name>/sim/files', methods=['GET'])
def list_sim_files(study_name: str):
    """
    List all files in the study's sim/ folder.

    Args:
        study_name: Name of the study
    """
    try:
        study_dir = project_root / 'optimization_results' / study_name
        sim_dir = study_dir / 'sim'

        if not sim_dir.exists():
            return jsonify({
                'success': False,
                'error': f'Study {study_name} does not exist'
            }), 404

        # List all files
        files = []
        for file_path in sim_dir.iterdir():
            if file_path.is_file():
                files.append({
                    'name': file_path.name,
                    'size': file_path.stat().st_size,
                    'extension': file_path.suffix,
                    'modified': datetime.fromtimestamp(file_path.stat().st_mtime).isoformat()
                })

        # Check for .sim file
        has_sim = any(f['extension'] == '.sim' for f in files)

        return jsonify({
            'success': True,
            'files': files,
            'has_sim_file': has_sim,
            'sim_folder': str(sim_dir)
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/study/<study_name>/explore', methods=['POST'])
def explore_sim_file(study_name: str):
    """
    Explore the .sim file in the study folder to extract expressions.

    Args:
        study_name: Name of the study
    """
    try:
        study_dir = project_root / 'optimization_results' / study_name
        sim_dir = study_dir / 'sim'

        # Find .sim file
        sim_files = list(sim_dir.glob('*.sim'))
        if not sim_files:
            return jsonify({
                'success': False,
                'error': 'No .sim file found in sim/ folder'
            }), 404

        sim_file = sim_files[0]

        # Run NX journal to extract expressions
        import subprocess
        journal_script = project_root / 'dashboard' / 'scripts' / 'extract_expressions.py'
        output_file = study_dir / 'expressions.json'

        # Execute journal
        nx_executable = r"C:\Program Files\Siemens\Simcenter3D_2412\NXBIN\run_journal.exe"

        result = subprocess.run(
            [nx_executable, str(journal_script), str(sim_file), str(output_file)],
            capture_output=True,
            text=True,
            timeout=120
        )

        if result.returncode != 0:
            return jsonify({
                'success': False,
                'error': f'NX journal failed: {result.stderr}'
            }), 500

        # Load extracted expressions
        if not output_file.exists():
            return jsonify({
                'success': False,
                'error': 'Expression extraction failed - no output file'
            }), 500

        with open(output_file, 'r') as f:
            expressions = json.load(f)

        return jsonify({
            'success': True,
            'sim_file': str(sim_file),
            'expressions': expressions
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/study/<study_name>/config', methods=['GET'])
def get_study_config(study_name: str):
    """
    Get the configuration for a study.

    Args:
        study_name: Name of the study
    """
    try:
        study_dir = project_root / 'optimization_results' / study_name
        config_path = study_dir / 'config.json'

        if not config_path.exists():
            return jsonify({
                'success': True,
                'config': None,
                'message': 'No configuration found for this study'
            })

        with open(config_path, 'r') as f:
            config = json.load(f)

        return jsonify({
            'success': True,
            'config': config
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


@app.route('/api/study/<study_name>/config', methods=['POST'])
def save_study_config(study_name: str):
    """
    Save configuration for a study.

    Args:
        study_name: Name of the study

    Request body:
        {
            "design_variables": [...],
            "objectives": [...],
            "constraints": [...],
            "optimization_settings": {...}
        }
    """
    try:
        study_dir = project_root / 'optimization_results' / study_name

        if not study_dir.exists():
            return jsonify({
                'success': False,
                'error': f'Study {study_name} does not exist'
            }), 404

        config = request.get_json()
        config_path = study_dir / 'config.json'

        # Save configuration
        with open(config_path, 'w') as f:
            json.dump(config, f, indent=2)

        # Update metadata
        metadata_path = study_dir / 'metadata.json'
        if metadata_path.exists():
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)

            metadata['is_configured'] = True
            metadata['last_modified'] = datetime.now().isoformat()

            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=2)

        return jsonify({
            'success': True,
            'message': f'Configuration saved for study {study_name}'
        })

    except Exception as e:
        return jsonify({
            'success': False,
            'error': str(e)
        }), 500


if __name__ == '__main__':
    print("="*60)
    print("ATOMIZER DASHBOARD API")
    print("="*60)
    print("Starting Flask server on http://localhost:8080")
    print("Access the dashboard at: http://localhost:8080")
    print("="*60)
    app.run(debug=True, host='0.0.0.0', port=8080, threaded=True)