import os
import json
import yaml
from flask import Flask, request, jsonify, render_template
from pathlib import Path

app = Flask(__name__)

# Basic path configuration
CHARACTER_SCENE_POOL_PATH = "synthetic_character_scene_pool"
WORLDVIEW_PATH = "cache/worldview"
TRACKING_PATH = "cache/tracking.json"

# Store user selected file paths
selected_files = []
selected_worldview = None

@app.route('/')
def index():
    """Render main page"""
    return render_template('preview.html')

@app.route('/api/files')
def get_files():
    """Recursively get all files and directories in the character_pool directory"""
    return jsonify(scan_directory(CHARACTER_SCENE_POOL_PATH))

@app.route('/api/worldviews')
def get_worldviews():
    """Get all YAML files in the worldview directory"""
    return jsonify(scan_directory(WORLDVIEW_PATH))

@app.route('/api/tracking')
def get_tracking():
    """Get the content of tracking.json and add worldview name"""
    try:
        if os.path.exists(TRACKING_PATH):
            with open(TRACKING_PATH, 'r', encoding='utf-8') as f:
                tracking_data = json.load(f)
            
            # Add worldview name to each record
            for record in tracking_data:
                worldview_path = record['from_worldview']
                try:
                    with open(worldview_path, 'r', encoding='utf-8') as f:
                        worldview_content = yaml.safe_load(f)
                        record['worldview_name'] = worldview_content.get('Name', '')
                except Exception as e:
                    print(f"Read worldview file failed: {str(e)}")
                    record['worldview_name'] = ''
            
            return jsonify(tracking_data)
        else:
            return jsonify([])
    except Exception as e:
        return jsonify({"error": f"读取tracking.json失败: {str(e)}"}), 500

@app.route('/api/separated-characters')
def get_separated_characters():
    """Get character files in the specified directory"""
    path = request.args.get('path')
    if not path:
        return jsonify({"error": "No directory path provided"}), 400
    
    try:
        directory = Path(path)
        if not directory.exists() or not directory.is_dir():
            return jsonify({"error": "Directory does not exist"}), 404
        
        # Get all YAML files in the directory
        character_files = []
        for file_path in directory.glob("*.yaml"):
            character_files.append({
                "name": file_path.stem,
                "path": str(file_path),
                "type": "character"
            })
        
        return jsonify(character_files)
    except Exception as e:
        return jsonify({"error": f"Read character file failed: {str(e)}"}), 500

def scan_directory(directory):
    """Recursively scan the directory, return the file tree structure"""
    path = Path(directory)
    if not path.exists():
        # Try to create directory
        try:
            path.mkdir(parents=True, exist_ok=True)
            print(f"Create directory: {directory}")
        except Exception as e:
            return {"error": f"Directory {directory} does not exist and cannot be created: {str(e)}"}
    
    result = {
        "name": path.name or path.as_posix(),
        "type": "directory",
        "path": path.as_posix(),
        "children": []
    }
    
    try:
        # Get all files and subdirectories in the directory
        items = list(path.iterdir())
        items.sort(key=lambda x: (not x.is_dir(), x.name))  # Directory first, same type sorted by name
        
        for item in items:
            if item.is_dir():
                # Recursively process subdirectories
                result["children"].append(scan_directory(item))
            else:
                # Process file
                result["children"].append({
                    "name": item.name,
                    "type": "file",
                    "path": item.as_posix()
                })
        
        return result
    except Exception as e:
        return {"error": f"Scan directory {directory} failed: {str(e)}"}

@app.route('/api/file-content')
def get_file_content():
    """Get file content"""
    file_path = request.args.get('path')
    if not file_path:
        return "No file path provided", 400
    
    # 确保路径安全，防止目录遍历攻击
    path = Path(file_path)
    if not path.exists() or not path.is_file():
        return f"File {file_path} does not exist", 404
    
    # 确保文件在允许的目录下
    allowed_paths = [CHARACTER_SCENE_POOL_PATH, WORLDVIEW_PATH]
    if not any(str(path).startswith(allowed_path) for allowed_path in allowed_paths):
        return "Only files in the character pool or worldview directory can be accessed", 403
    
    try:
        # Read YAML file content
        with open(path, 'r', encoding='utf-8') as f:
            content = f.read()
            
        # If it is a YAML file, try to format the content
        if path.suffix.lower() in ['.yaml', '.yml']:
            try:
                yaml_content = yaml.safe_load(content)
                # Here you can process the YAML content, such as formatting display
                return content
            except Exception as e:
                return f"Parse YAML file failed: {str(e)}"
        
        return content
    except Exception as e:
        return f"Read file failed: {str(e)}", 500

@app.route('/api/confirm-selection', methods=['POST'])
def confirm_selection():
    """Confirm selected files"""
    global selected_files, selected_worldview
    
    data = request.json
    if not data:
        return jsonify({"success": False, "error": "No selection data provided"}), 400
    
    # Check character selection
    if 'selected_files' not in data:
        return jsonify({"success": False, "error": "No selected character list provided"}), 400
    
    # Check worldview selection
    if 'selected_worldview' not in data:
        return jsonify({"success": False, "error": "No selected worldview provided"}), 400
    
    selected_files = data['selected_files']
    selected_worldview = data['selected_worldview']
    
    # Verify character number
    if len(selected_files) < 2:
        return jsonify({
            "success": False, 
            "error": "At least two characters need to be selected"
        }), 400
    
    # Verify all selected files exist
    valid_files = []
    invalid_files = []
    
    for file_path in selected_files:
        path = Path(file_path)
        if path.exists() and path.is_file():
            valid_files.append(file_path)
        else:
            invalid_files.append(file_path)
    
    # Verify worldview file exists
    if selected_worldview:
        worldview_path = Path(selected_worldview)
        if not (worldview_path.exists() and worldview_path.is_file()):
            return jsonify({
                "success": False, 
                "error": f"Selected worldview file does not exist: {selected_worldview}"
            }), 400
    else:
        return jsonify({
            "success": False, 
            "error": "Must select a worldview"
        }), 400
    
    # Update selected files to valid files
    selected_files = valid_files
    
    result = {
        "success": True,
        "selected_files": selected_files,
        "selected_worldview": selected_worldview
    }
    
    if invalid_files:
        result["warnings"] = f"The following files do not exist or are invalid: {', '.join(invalid_files)}"
    
    # Save selected files to file, for other programs to use
    save_selection_to_file()
    
    return jsonify(result)

def save_selection_to_file():
    """Save selected files to file"""
    try:
        with open('../selected_characters.json', 'w', encoding='utf-8') as f:
            json.dump({
                "selected_files": selected_files,
                "selected_worldview": selected_worldview
            }, f, ensure_ascii=False, indent=2)
    except Exception as e:
        print(f"Save selected files failed: {e}")

@app.route('/api/get-selected-files')
def get_selected_files():
    """Get current selected file list"""
    global selected_files, selected_worldview
    return jsonify({
        "selected_files": selected_files,
        "selected_worldview": selected_worldview
    })

@app.route('/api/use-selection-for-dialogue')
def use_selection_for_dialogue():
    """Use selected files for dialogue generation"""
    global selected_files, selected_worldview
    
    if not selected_files:
        return jsonify({"success": False, "error": "No character files selected"}), 400
    
    if not selected_worldview:
        return jsonify({"success": False, "error": "No worldview file selected"}), 400
    
    # Build command line arguments for calling generate_dialogue
    character_args = " ".join([f'"{f}"' for f in selected_files])
    command = f"python initial_dialogue_generation.py --worldview \"{selected_worldview}\" --characters {character_args}"
    
    # Here you can actually call the subprocess to execute the command
    # But here only return the command, user can execute it by themselves
    return jsonify({
        "success": True,
        "command": command,
        "files": selected_files,
        "worldview": selected_worldview
    })

@app.route('/api/current-selection')
def get_current_selection():
    """Get current selected files and worldview"""
    global selected_files, selected_worldview
    return jsonify({
        "selected_files": selected_files,
        "selected_worldview": selected_worldview
    })

@app.route('/api/character-total')
def get_character_total():
    """Get total number of characters in the character pool"""
    try:
        total_count = 0
        pool_path = Path(CHARACTER_SCENE_POOL_PATH)
        
        # Recursively traverse all .yaml files
        for yaml_file in pool_path.rglob("*.yaml"):
            total_count += 1
            
        return jsonify({
            "success": True,
            "total_count": total_count
        })
    except Exception as e:
        return jsonify({
            "success": False,
            "error": f"Count character total failed: {str(e)}"
        }), 500

if __name__ == '__main__':
    # Ensure necessary directories exist
    required_dirs = [CHARACTER_SCENE_POOL_PATH, WORLDVIEW_PATH]
    for directory in required_dirs:
        if not os.path.exists(directory):
            os.makedirs(directory)
            print(f"Create directory: {directory}")
    
    # Start server
    app.run(debug=True, host='0.0.0.0', port=1210)
