main.py

import argparse
import logging
import os
import re
import subprocess
import sys
import threading
import time
import pyperclip
sys.path.insert(0, os.path.dirname(__file__))

from core.brain import Brain
from core.listener import Listener
from core.speaker import Speaker
from modules.automator import Automator
from modules.system import System
from modules.web_search import WebSearch
from modules.vision import Vision
from modules.hand_control import HandControl
from modules.app_control import AppController
from modules.security import Security
from config.logger import setup_logging
from config.settings import WAKE_WORD, IDLE_TIMEOUT, VERSION, APP_NAME

try:
    import keyboard as kb
    _HAS_KB = True
except ImportError:
    _HAS_KB = False


def _show_toast(title: str, message: str):
    ps = (
        '[Windows.UI.Notifications.ToastNotificationManager,'
        ' Windows.UI.Notifications, ContentType = WindowsRuntime] > $null;'
        '$t = [Windows.UI.Notifications.ToastNotificationManager]::'
        'GetTemplateContent([Windows.UI.Notifications.ToastTemplateType]::ToastText02);'
        '$n = $t.GetElementsByTagName("text");'
        f'$n.Item(0).AppendChild($t.CreateTextNode("{title}")) > $null;'
        f'$n.Item(1).AppendChild($t.CreateTextNode("{message}")) > $null;'
        '$toast = [Windows.UI.Notifications.ToastNotification]::new($t);'
        f'[Windows.UI.Notifications.ToastNotificationManager]::'
        f'CreateToastNotifier("{APP_NAME}").Show($toast)'
    )
    try:
        subprocess.run(
            ["powershell", "-NoProfile", "-Command", ps],
            capture_output=True, timeout=10,
            creationflags=subprocess.CREATE_NO_WINDOW,
        )
    except Exception:
        pass


def _install_startup():
    startup = os.path.join(
        os.environ["APPDATA"],
        "Microsoft", "Windows", "Start Menu", "Programs", "Startup",
    )
    shortcut = os.path.join(startup, "Jarvis.lnk")
    target = sys.executable
    script = os.path.abspath(__file__)
    ps = (
        '$ws = New-Object -ComObject WScript.Shell;'
        f'$s = $ws.CreateShortcut("{shortcut}");'
        f'$s.TargetPath = "{target}";'
        f'$s.Arguments = "main.py --startup";'
        f'$s.WorkingDirectory = "{os.path.dirname(script)}";'
        '$s.Description = "Jarvis PC Assistant";'
        '$s.Save()'
    )
    subprocess.run(
        ["powershell", "-NoProfile", "-Command", ps],
        check=True, creationflags=subprocess.CREATE_NO_WINDOW,
    )
    print(f"Startup shortcut created: {shortcut}")
    print("Jarvis will now start automatically when you log in.")


def _uninstall_startup():
    startup = os.path.join(
        os.environ["APPDATA"],
        "Microsoft", "Windows", "Start Menu", "Programs", "Startup",
    )
    shortcut = os.path.join(startup, "Jarvis.lnk")
    if os.path.exists(shortcut):
        os.remove(shortcut)
        print(f"Removed startup shortcut: {shortcut}")
    else:
        print("No startup shortcut found.")


def _check_environment():
    logger = logging.getLogger("jarvis")
    try:
        import ollama
        client = ollama.Client(host="http://localhost:11434")
        client.list()
        logger.info("Ollama connection OK")
    except Exception:
        logger.warning(
            "Ollama not reachable. Install from https://ollama.com "
            "and run 'ollama serve'. Brain features will be unavailable."
        )
    try:
        import speech_recognition as sr
        with sr.Microphone():
            pass
        logger.info("Microphone OK")
    except Exception as e:
        logger.warning(f"Microphone unavailable: {e}")
    try:
        import cv2
        cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
        if cap.isOpened():
            cap.release()
            logger.info("Camera OK")
    except Exception:
        pass


class Jarvis:
    def __init__(self):
        self.logger = logging.getLogger("jarvis")
        self.brain = Brain()
        self.listener = Listener()
        self.speaker = Speaker()
        self.automator = Automator()
        self.system = System()
        self.web = WebSearch()
        self.vision = Vision()
        self.hand = HandControl()
        self.apps = AppController()
        self.running = True
        self._timers = []

    _power_map = {
        "shutdown": "shutdown /s /t 3",
        "restart": "shutdown /r /t 3",
        "hibernate": "shutdown /h",
        "sleep": "rundll32.exe powrprof.dll,SetSuspendState 0,1,0",
    }

    def _start_timer(self, seconds: int, label: str = ""):
        t = threading.Timer(seconds, self._timer_done, args=[label])
        t.daemon = True
        t.start()
        self._timers.append(t)

    def _timer_done(self, label: str):
        import winsound
        for freq, dur in [(1000, 200), (1200, 200), (1500, 400)]:
            winsound.Beep(freq, dur)
        msg = f"{label} done!" if label else "Timer done!"
        self.speaker.say(msg)
        _show_toast("Timer Done", msg)

    def execute_action(self, action: dict) -> str:
        action_type = action.get("action")
        self.logger.debug("Action: %s", action)
        if action_type == "type":
            self.automator.type_text(action.get("text", ""))
        elif action_type == "press":
            self.automator.press_key(action.get("key", ""))
        elif action_type == "keypress":
            self.automator.press_key(action.get("key", ""))
        elif action_type == "hotkey":
            self.automator.hotkey(action.get("keys", []))
        elif action_type == "click":
            self.automator.click(action.get("x"), action.get("y"), action.get("button", "left"))
        elif action_type == "double_click":
            self.automator.double_click(action.get("x"), action.get("y"))
        elif action_type == "right_click":
            self.automator.right_click(action.get("x"), action.get("y"))
        elif action_type == "move_mouse":
            self.automator.move_mouse(action["x"], action["y"])
        elif action_type == "scroll":
            self.automator.scroll(action.get("clicks", 0))
        elif action_type == "screenshot":
            path = self.automator.screenshot()
            return f"Screenshot saved to {path}."
        elif action_type == "run":
            return self.system.run_command(action.get("command", ""))
        elif action_type == "open":
            return self.system.open_path(action.get("path", action.get("command", "")))
        elif action_type == "write_text":
            self.automator.write_slowly(action.get("text", ""))
        elif action_type == "system_info":
            return self.system.get_system_info()
        elif action_type == "battery":
            return self.system.get_battery()
        elif action_type == "time":
            from datetime import datetime
            return datetime.now().strftime("%I:%M %p").lstrip("0")
        elif action_type == "date":
            from datetime import datetime
            return datetime.now().strftime("%A, %B %d, %Y")
        elif action_type == "describe_screen":
            return self.vision.describe()
        elif action_type == "read_screen":
            return self.vision.read_text()
        elif action_type == "click_on":
            return self.vision.click_on(action.get("target", ""))
        elif action_type == "find_and_type":
            return self.vision.find_and_type(action.get("target", ""), action.get("text", ""))
        elif action_type == "hand_control":
            if action.get("state") == "on":
                self.hand.start()
                return "Hand control activated."
            elif action.get("state") == "off":
                self.hand.stop()
                return "Hand control deactivated."
            else:
                self.hand.toggle()
                return f"Hand control {'on' if self.hand.active else 'off'}."

        elif action_type == "volume_up" and _HAS_KB:
            kb.press_and_release("volume up")
        elif action_type == "volume_down" and _HAS_KB:
            kb.press_and_release("volume down")
        elif action_type == "volume_mute" and _HAS_KB:
            kb.press_and_release("volume mute")
        elif action_type == "media_play_pause" and _HAS_KB:
            kb.press_and_release("play/pause media")
        elif action_type == "media_next" and _HAS_KB:
            kb.press_and_release("next track")
        elif action_type == "media_prev" and _HAS_KB:
            kb.press_and_release("previous track")

        elif action_type == "copy":
            self.automator.hotkey(["ctrl", "c"])
        elif action_type == "cut":
            self.automator.hotkey(["ctrl", "x"])
        elif action_type == "paste":
            self.automator.hotkey(["ctrl", "v"])
        elif action_type == "read_clipboard":
            return f"Clipboard: {pyperclip.paste()}"

        elif action_type == "show_desktop":
            self.automator.hotkey(["win", "d"])
        elif action_type == "maximize_window":
            self.automator.hotkey(["win", "up"])
        elif action_type == "minimize_window":
            self.automator.hotkey(["win", "down"])
        elif action_type == "close_window":
            self.automator.hotkey(["alt", "f4"])
        elif action_type == "switch_window":
            self.automator.hotkey(["alt", "tab"])

        elif action_type == "lock":
            self.automator.hotkey(["win", "l"])
        elif action_type in self._power_map:
            return self.system.run_command(self._power_map[action_type])

        elif action_type == "list_processes":
            procs = self.system.list_processes()
            names = list(dict.fromkeys(p["name"] for p in procs[:15] if p["name"]))
            return f"Running: {', '.join(names)}."

        elif action_type == "timer":
            seconds = action.get("seconds", 0)
            label = action.get("label", "")
            self._start_timer(seconds, label)
            return f"Timer set for {seconds} seconds."

        else:
            return f"Unknown action: {action_type}"
        return f"Executed {action_type}."

    _app_map = {
        "calculator": "calc",
        "notepad": "notepad",
        "chrome": "chrome",
        "google chrome": "chrome",
        "edge": "msedge",
        "firefox": "firefox",
        "code": "code",
        "vscode": "code",
        "visual studio code": "code",
        "cmd": "cmd",
        "command prompt": "cmd",
        "powershell": "powershell",
        "explorer": "explorer",
        "file explorer": "explorer",
        "slack": "slack",
        "discord": "discord",
        "spotify": "spotify",
    }

    _exe_map = {
        "chrome": "chrome.exe",
        "calculator": "calc.exe",
        "notepad": "notepad.exe",
        "cmd": "cmd.exe",
        "powershell": "powershell.exe",
        "explorer": "explorer.exe",
        "edge": "msedge.exe",
        "firefox": "firefox.exe",
        "code": "Code.exe",
        "slack": "slack.exe",
        "discord": "discord.exe",
        "spotify": "Spotify.exe",
    }

    def _resolve_app(self, name: str) -> str:
        return self._app_map.get(name, name)

    def _fast_path(self, text: str) -> tuple[str, dict] | None:
        t = text.lower().strip()

        known_commands = {
            "exit", "quit", "goodbye", "shutdown", "restart", "reboot",
            "hibernate", "sleep",
            "hello", "hi", "hey",
        }

        if t in ("open hand control",):
            return "Opening hand control.", {"action": "hand_control", "state": "on"}

        if t in ("volume up", "turn it up", "increase volume", "louder"):
            return "Turning volume up.", {"action": "volume_up"}
        if t in ("volume down", "turn it down", "decrease volume", "quieter"):
            return "Turning volume down.", {"action": "volume_down"}
        if t in ("mute", "volume mute", "silence", "turn off sound"):
            return "Muting.", {"action": "volume_mute"}
        if t in ("unmute", "turn on sound"):
            return "Unmuting.", {"action": "volume_mute"}

        if t in ("play", "pause", "play pause", "play/pause", "resume"):
            return "Play pause.", {"action": "media_play_pause"}
        if t in ("next track", "next song", "skip", "next"):
            return "Next track.", {"action": "media_next"}
        if t in ("previous track", "previous song", "go back", "previous"):
            return "Previous track.", {"action": "media_prev"}

        if t in ("copy", "copy that"):
            return "Copying.", {"action": "copy"}
        if t in ("cut", "cut that"):
            return "Cutting.", {"action": "cut"}
        if t in ("paste", "paste that"):
            return "Pasting.", {"action": "paste"}
        clip_match = re.match(r"^(?:what'?s|what is|read|show)\s+(?:on\s+)?(?:my\s+)?clipboard$", t)
        if clip_match:
            return "Reading clipboard.", {"action": "read_clipboard"}

        if t in ("show desktop", "minimize all", "show all windows"):
            return "Showing desktop.", {"action": "show_desktop"}
        if t in ("maximize", "maximize window", "full screen", "fullscreen", "max window"):
            return "Maximizing window.", {"action": "maximize_window"}
        if t in ("minimize", "minimize window", "min window"):
            return "Minimizing window.", {"action": "minimize_window"}
        if t in ("close window", "close this", "close window please"):
            return "Closing window.", {"action": "close_window"}
        if t in ("switch window", "switch app", "alt tab", "switch windows"):
            return "Switching window.", {"action": "switch_window"}

        if t in ("lock", "lock pc", "lock computer", "lock screen"):
            return "Locking.", {"action": "lock"}

        timer_match = re.match(r"^(?:set\s+)?(?:a\s+)?timer\s+(?:for\s+)?(\d+)\s*(s|sec|secs|seconds?|m|min|mins|minutes?)?\s*$", t)
        if timer_match:
            value = int(timer_match.group(1))
            unit = (timer_match.group(2) or "s").lower()
            seconds = value * 60 if unit.startswith("m") else value
            return f"Timer set for {seconds} seconds.", {"action": "timer", "seconds": seconds, "label": f"{value} {'min' if unit.startswith('m') else 'sec'}"}

        if t in ("list processes", "running programs", "whats running", "what is running"):
            return "Listing processes.", {"action": "list_processes"}

        open_match = re.match(r"^open\s+(.+)$", t)
        if open_match:
            app = self._resolve_app(open_match.group(1))
            return f"Opening {app}.", {"action": "run", "command": app}

        search_on_match = re.match(r"^search\s+(?:up\s+)?(.+?)\s+on\s+(.+)$", t)
        if search_on_match:
            query, app = search_on_match.group(1), search_on_match.group(2)
            return f"Searching {app} for {query}.", {"action": "app_search", "app": app, "query": query}

        search_match = re.match(r"^search\s+(?:up\s+)?(.+)$", t)
        if search_match:
            query = search_match.group(1)
            return f"Searching for {query}.", {"action": "app_search", "query": query}

        look_match = re.match(r"^(?:look up|find)\s+(.+?)(?:\s+on\s+(.+))?$", t)
        if look_match:
            query = look_match.group(1)
            app = look_match.group(2)
            d = {"action": "app_search", "query": query}
            if app:
                d["app"] = app
            return f"Looking up {query}.", d

        if t in ("time", "what time is it", "what's the time", "tell me the time"):
            from datetime import datetime
            return datetime.now().strftime("It's %I:%M %p.").lstrip("0"), {"action": "time"}

        if t in ("date", "what is the date", "what's the date", "today's date", "todays date"):
            from datetime import datetime
            return datetime.now().strftime("It's %A, %B %d, %Y."), {"action": "date"}

        if t in ("screenshot", "take a screenshot"):
            return "Taking screenshot.", {"action": "screenshot"}

        if t in ("system info", "system_info", "system information", "pc info"):
            return "Getting system info.", {"action": "system_info"}

        if t in ("battery", "battery status", "battery level"):
            return "Checking battery.", {"action": "battery"}

        type_match = re.match(r"^type\s+(.+)$", t)
        if type_match:
            return f"Typing {type_match.group(1)}.", {"action": "type", "text": type_match.group(1)}

        press_match = re.match(r"^press\s+(.+)$", t)
        if press_match:
            return f"Pressing {press_match.group(1)}.", {"action": "press", "key": press_match.group(1)}

        if t in ("click", "click here", "click there"):
            return "Clicking.", {"action": "click"}

        close_match = re.match(r"^close\s+(.+)$", t)
        if close_match:
            app = close_match.group(1)
            exe = self._exe_map.get(app, f"{app}.exe")
            return f"Closing {app}.", {"action": "run", "command": f"taskkill /f /im {exe}"}

        hand_match = re.match(r"^(turn\s+)?(on|off)\s+hand\s+control$", t)
        if hand_match:
            state = hand_match.group(2)
            return f"Turning {state} hand control.", {"action": "hand_control", "state": state}

        if t in ("hand control",):
            return None

        if t in self._app_map:
            app = self._resolve_app(t)
            return f"Opening {app}.", {"action": "run", "command": app}

        if t in known_commands:
            return None

        words = t.split()
        if len(words) == 1:
            return f"Searching {words[0]}.", {"action": "app_search", "query": words[0]}

        return None

    def _confirm_dangerous(self, description: str) -> bool:
        self.speaker.say(description)
        response = self.listener.listen(timeout=5)
        if response and any(w in response for w in ("yes", "yeah", "sure", "yep", "confirm", "do it", "go ahead")):
            self.speaker.say("Okay.")
            return True
        self.speaker.say("Cancelled.")
        return False

    def process_text(self, text: str) -> str | None:
        if not text:
            return None
        self.logger.info("You: %s", text)

        fast = self._fast_path(text)
        from_fast = fast is not None
        if fast:
            response_text, action = fast
            self.logger.info("Jarvis: %s", response_text)
            self.speaker.say(response_text)
        else:
            response_text, action = self.brain.think(text)
            if response_text:
                self.logger.info("Jarvis: %s", response_text)
                self.speaker.say(response_text)

        if action:
            actions = action if isinstance(action, list) else [action]
            for act in actions:
                action_type = act.get("action")
                level = Security.classify_action(action_type, act)
                if level == "blocked":
                    msg = "I cannot do that - it is blocked for security reasons."
                    self.logger.warning(msg)
                    self.speaker.say(msg)
                    continue
                if level == "dangerous":
                    descs = {
                        "shutdown": "Shutting down can lose unsaved work.",
                        "restart": "Restarting can lose unsaved work.",
                        "hibernate": "Hibernating the system.",
                        "sleep": "Putting the system to sleep.",
                    }
                    desc = descs.get(action_type, "This operation is potentially dangerous.")
                    if not self._confirm_dangerous(f"{desc} Are you sure?"):
                        continue
                if action_type in ("run", "open"):
                    tracked = act.get("command") or act.get("path") or ""
                    self.apps.set_last_app(tracked)
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                    if result and "Could not find" in result:
                        query = tracked
                        fallback_result = self.apps.search(query)
                        self.logger.info("Fell back to searching '%s' in active window.", query)
                        self.logger.info("[%s]", fallback_result)
                elif action_type == "web_search":
                    query = act.get("query", text)
                    self.logger.debug("Searching web: %s", query)
                    results = self.web.search(query)
                    self.logger.debug("Got search results")
                    final_text, _ = self.brain.think(f"Based on these search results, answer my question briefly.\n\nSearch results:\n{results}")
                    if final_text:
                        self.logger.info("Jarvis: %s", final_text)
                        self.speaker.say(final_text)
                    response_text = final_text or response_text
                elif action_type == "app_search":
                    query = act.get("query", text)
                    app = act.get("app")
                    result = self.apps.search(query, app)
                    self.logger.info("[%s]", result)
                elif action_type in ("describe_screen", "read_screen"):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("Result: %s", result)
                        self.speaker.say(result)
                        response_text = result
                elif action_type == "hand_control":
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        self.speaker.say(result)
                elif action_type in ("time", "date"):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        if not from_fast:
                            self.speaker.say(result)
                        response_text = result
                elif action_type in ("click_on", "find_and_type"):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        if "Could not find" in result:
                            self.speaker.say(f"I could not find {act.get('target', 'it')} on screen.")
                elif action_type in ("read_clipboard", "list_processes", "system_info", "battery"):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        self.speaker.say(result)
                        response_text = result
                elif action_type in ("timer",):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        if not from_fast:
                            self.speaker.say(result)
                elif action_type.startswith(("volume_", "media_", "show_desktop", "maximize_", "minimize_", "close_window", "switch_window")):
                    result = self.execute_action(act)
                    if result:
                        self.logger.debug("[%s]", result)
                elif action_type in ("copy", "cut", "paste"):
                    self.execute_action(act)
                elif action_type == "lock":
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        self.speaker.say("Locking.")
                elif action_type in ("shutdown", "restart", "hibernate", "sleep"):
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
                        power_labels = {"shutdown": "Shutting down", "restart": "Restarting", "hibernate": "Hibernating", "sleep": "Sleeping"}
                        self.speaker.say(power_labels.get(action_type, action_type.capitalize()) + ".")
                else:
                    result = self.execute_action(act)
                    if result:
                        self.logger.info("[%s]", result)
        return response_text

    def run(self):
        self.logger.info("Initializing")
        self.listener.clap.start_monitor()
        self.speaker.say("Jarvis online. Clap to activate.")
        while self.running:
            try:
                if not self.listener.clap.clap_detected:
                    time.sleep(0.05)
                    continue
                self.logger.info("Clap detected")
                self.speaker.say("Yes?")
                time.sleep(0.4)
                active = True
                idle_since = time.time()
                while active and self.running:
                    command = self.listener.listen(timeout=3)
                    if command:
                        idle_since = time.time()
                        if self.listener.wake_word in command:
                            self.logger.info("Deactivated")
                            active = False
                            self.listener.clap.reset_clap()
                            self.speaker.say("Going to sleep.")
                            break
                        self.logger.info("You said: %s", command)
                        if command in ("exit", "quit", "goodbye"):
                            self.speaker.say("Goodbye.")
                            self.running = False
                            break
                        self.process_text(command)
                    elif time.time() - idle_since > IDLE_TIMEOUT:
                        self.logger.info("Idle timeout")
                        self.listener.clap.reset_clap()
                        self.speaker.say("Going to sleep.")
                        break
                    else:
                        self.logger.info("Listening...")
            except KeyboardInterrupt:
                self.logger.info("Goodbye")
                self.listener.clap.stop_monitor()
                break
        self.listener.clap.stop_monitor()

    def run_cli(self):
        self.logger.info("CLI mode. Type 'exit' to quit.")
        while self.running:
            try:
                user_input = input("  You: ").strip()
                if user_input in ("exit", "quit"):
                    self.logger.info("Goodbye")
                    break
                if user_input:
                    self.process_text(user_input)
            except KeyboardInterrupt:
                self.logger.info("Goodbye")
                break


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=f"{APP_NAME} PC Assistant v{VERSION}")
    parser.add_argument("--cli", action="store_true", help="Run in text-only CLI mode")
    parser.add_argument("--startup", action="store_true", help="Run at Windows startup (adds delay)")
    parser.add_argument("--install", action="store_true", help="Add Jarvis to Windows startup")
    parser.add_argument("--uninstall", action="store_true", help="Remove Jarvis from Windows startup")
    parser.add_argument("--version", action="store_true", help="Show version and exit")
    args = parser.parse_args()

    if args.version:
        print(f"{APP_NAME} v{VERSION}")
        sys.exit(0)

    setup_logging()

    if args.install:
        _install_startup()
        sys.exit(0)
    if args.uninstall:
        _uninstall_startup()
        sys.exit(0)

    if args.startup:
        logger = logging.getLogger("jarvis")
        logger.info("Startup launch, waiting 5s for system ready...")
        time.sleep(5)

    if not args.cli:
        _check_environment()

    app = Jarvis()
    if args.cli:
        app.run_cli()
    else:
        app.run()