import os
import sys

from PytorchBoot.application import application_class
from PytorchBoot.stereotype import get_all_component_classes, get_all_component_comments
from PytorchBoot.utils.log_util import Log
from PytorchBoot.utils.timer_util import Timer
from PytorchBoot.utils.project_util import ProjectUtil
from PytorchBoot.templates.application import template as app_template
from PytorchBoot.templates.config import template as config_template
from PytorchBoot.ui.server.app import app


def run():
    root_path = os.getcwd()
    ProjectUtil.scan_project(root_path)
    
    app_name = "default"
    if len(application_class) == 0:
        Log.error("No class annotated with @PytorchBootApplication found.", True)
    if len(sys.argv) < 3 and "default" not in application_class:
        Log.error("No default @PytorchBootApplication found. Please specify the 'name' parameter.", True)
    if len(sys.argv) == 3:
        app_name = sys.argv[2]
        
    app_cls = application_class.get(app_name)
        
    if app_cls is None:
        Log.error(f"No class annotated with @PytorchBootApplication found with the name '{app_name}'.", True)
    
    if not hasattr(app_cls, "start"):
        Log.error("The class annotated with @PytorchBootApplication should have a 'start' method.", True)
    
    Log.info(f"Application '{app_cls.__name__}' started.")
    timer = Timer("Application")

    timer.start()
    app_cls.start()
    timer.stop()
    Log.info(timer.get_elasped_time_str(Timer.HOURS))
    Log.success("Application finished.")


def init():
    Log.info("Initializing PytorchBoot project.")
    root_path = os.getcwd()
    if len(os.listdir(root_path)) > 0:
        Log.error("Current directory is not empty. Please provide an empty directory.")
    else:
        with open(os.path.join(root_path, "application.py"), "w") as file:
            file.write(app_template)
        with open(os.path.join(root_path, "config.yaml"), "w") as file:
            file.write(config_template)
            
        Log.success("PytorchBoot project initialized.")
        Log.info("Now you can create your components and run the application.")

def scan():
    root_path = os.getcwd()
    ProjectUtil.scan_project(root_path)
    comments = get_all_component_comments()
    Log.info("Components detected in the project:")
    for stereotype, classes in get_all_component_classes().items():
        Log.info(f"  {stereotype}:")
        for name, cls in classes.items():
            comment = comments[stereotype].get(name)
            if comment is not None:
                Log.warning(f"    - {name}: {cls.__module__}.{cls.__name__} ({comment})")
            else:
                Log.success(f"    - {name}: {cls.__module__}.{cls.__name__}")
        
    Log.info("Applications detected in the project:")
    for app_name, app_cls in application_class.items():
        Log.success(f"    - {app_name}: {app_cls.__module__}.{app_cls.__name__}")
    Log.success("Scan completed.")
    
def ui():
    port = 5000
    if len(sys.argv) == 3:
        port = int(sys.argv[2])
    Log.success(f"PytorchBoot UI server started at http://localhost:{port}")
    app.run(port=port, host="0.0.0.0")
    
    
def help():
    Log.info("PytorchBoot commands:")
    Log.info("  init: Initialize a new PytorchBoot project in the current directory.")
    Log.info("  run [name]: Run the PytorchBoot application with the specified name. If no name is provided, the default application will be run.")
    Log.info("  scan: Scan the project for PytorchBoot components.")
    Log.info("  ui [port]: Start the PytorchBoot UI server. If no port is provided, the default port 5000 will be used.")
    Log.info("  help: Display this help message.")

def main():
    if len(sys.argv) > 1:
        if sys.argv[1] == "init":
            init()
        elif sys.argv[1] == "run":
            run()
        elif sys.argv[1] == "scan":
            scan()
        elif sys.argv[1] == "ui":
            ui()
        elif sys.argv[1] == "help":
            help()
        else:
            Log.error("Invalid command: " + sys.argv[1] + ". Use 'pytorch-boot help' for help.")
    else:
        Log.error("Please provide a command to run the application.")