from hiveos.sim_agent import SimAgent
from hiveos.core_loop import ingest_intent, tick, core
from hiveos.db import get_db
from hiveos.utilities.sanitized_inputs import get_positive_int, get_non_empty_string, get_capability

import time
import yaml
import random

from hiveos.agents_sim.sim_display_agent import SimDisplayAgent
from hiveos.agents_sim.sim_terminal_agent import SimTerminalAgent
from hiveos.agents_sim.sim_sensor_agent import SimSensorAgent

from hiveos.agents_hardware.oled_agent import OledDisplayAgent
from hiveos.agents_hardware.piezo_agent import PiezoAgent

import threading
tick_loop_active = False
tick_thread = None
tick_interval = 0.1


def get_name_from_runtime(agent_id):
    runtime_agent = core._runtime_agents.get(agent_id)
    agent = runtime_agent.obj if runtime_agent else None
    return agent.name if agent else agent_id[:8]

def inject_task(name=None, count=None, cap=None):
    
    if name is None:
        name = get_non_empty_string("Task name: ")
        if name is None:
            return

    if count is None:
        count = get_positive_int("Chunks: ")
        if count is None:
            return

    if cap is None:
        cap = get_capability("Capability (default = 'generic'): ")

    chunk_list = []
    for i in range(count):
        chunk_list.append({
            "chunk_id": f"C{i+1}",
            "payload": f"Payload_{i+1}"
        })

    intent = {
        "task_name": name,
        "capability_required": cap,
        "execution_mode": "concurrent",
        "chunks": chunk_list
    }
    """
    intent = {
        "task_name": "wave_seq_test",
        "capability_required": "display",
        "execution_mode": "sequential",
        "chunks": [
            {"chunk_id": "C1", "payload": "Frame1"},
            {"chunk_id": "C2", "payload": "Frame2", "depends_on": "C1"},
            {"chunk_id": "C3", "payload": "Frame3", "depends_on": "C2"},
            {"chunk_id": "C4", "payload": "Frame4", "depends_on": "C3"},
            {"chunk_id": "C5", "payload": "Frame4", "depends_on": "C4"},
            {"chunk_id": "C6", "payload": "Frame4", "depends_on": "C5"},
            {"chunk_id": "C7", "payload": "Frame4", "depends_on": "C6"},
            {"chunk_id": "C8", "payload": "Frame4", "depends_on": "C7"},
            {"chunk_id": "C9", "payload": "Frame4", "depends_on": "C8"},
            {"chunk_id": "C10", "payload": "Frame4", "depends_on": "C9"},
            {"chunk_id": "C11", "payload": "Frame4", "depends_on": "C10"},
            {"chunk_id": "C12", "payload": "Frame4", "depends_on": "C11"},
            {"chunk_id": "C13", "payload": "Frame4", "depends_on": "C12"},
            {"chunk_id": "C14", "payload": "Frame4", "depends_on": "C13"},
            {"chunk_id": "C15", "payload": "Frame4", "depends_on": "C14"},
            {"chunk_id": "C16", "payload": "Frame4", "depends_on": "C15"},
            {"chunk_id": "C17", "payload": "Frame4", "depends_on": "C16"},
            {"chunk_id": "C18", "payload": "Frame4", "depends_on": "C17"},
            {"chunk_id": "C19", "payload": "Frame4", "depends_on": "C18"},
            {"chunk_id": "C20", "payload": "Frame4", "depends_on": "C19"},
        ]
    }
    """


    ingest_intent(intent)
    print(f"✅ Injected task '{name}' with {count} chunks.")


def inject_agents():
    count = get_positive_int("How many agents to create? ")
    if count is None:
        return
    cap = get_capability("Enter capability (default = 'generic'): ")
    for _ in range(count):
        core.inject_agent(SimAgent(capabilities=[cap]))
    print(f"✅ Injected {count} agent(s) with capability '{cap}'.")
    
def agent_shotgun():
    for i in range(3000):
        cap = random.choice(['display', 'oled', 'sensor', 'speaker', 'drone'])
        name = f"A{i}_{cap}"
        core.inject_agent(SimAgent(agent_id=name, name=name, capabilities=[cap]))
    
def inject_sim_display():
    new_agent = SimDisplayAgent(capability="display")
    core.inject_agent(new_agent)
    print(f"✅ SimDisplayAgent injected: {new_agent.agent_id}")
    
def inject_sim_terminal():
    new_agent = SimTerminalAgent()
    core.inject_agent(new_agent)
    print(f"✅ SimTerminalAgent injected: {new_agent.agent_id}")
    
def inject_sim_sensor():
    new_agent = SimSensorAgent()
    core.inject_agent(new_agent)
    print(f"✅ SimSensorAgent injected: {new_agent.agent_id}")
    
def inject_hw_oled():
    ble_address = "95:6F:67:9A:7A:F0"
    new_agent = OledDisplayAgent(ble_address=ble_address)
    core.inject_agent(new_agent)
    print(f"✅ HW Oled injected: {new_agent.agent_id}")
    
def inject_hw_piezo():
    ble_address = "59:C8:5F:EA:F9:C1"
    new_agent = PiezoAgent(ble_address=ble_address)
    core.inject_agent(new_agent)
    print(f"✅ HW Piezo injected: {new_agent.agent_id}")

def run_tick(verbose=False):
    print("\n🔄 Running system tick...")
    tick(verbose=verbose)
    print("✅ Tick complete.\n")
    print()
    
def background_tick_loop():
    global tick_loop_active
    while tick_loop_active:
        tick(verbose=False)
        time.sleep(tick_interval)  # Tick interval (can be adjusted)


def show_agents():
    with get_db() as conn:
        agents = conn.execute("SELECT * FROM agents").fetchall()
        print("\n📋 Registered Agents:")
        for row in agents:
            #ensure_registered(row)
            name = get_name_from_runtime(row[0])
            print(f"{name}: Status={row[1]}, Battery={row[2]}, Zone={row[3]}, {row[4]}")

def show_tasks():
    with get_db() as conn:
        tasks = conn.execute("SELECT * FROM tasks").fetchall()
        chunks = conn.execute("SELECT * FROM task_chunks").fetchall()

    chunk_map = {}
    for chunk in chunks:
        task_id = chunk[1]
        chunk_map.setdefault(task_id, []).append(chunk)

    print("\n📋 Tasks Overview:")
    for task in tasks:
        task_id, task_name = task[:2]
        runtime_task = core.task_objects.get(task_id)
        status = runtime_task.status if runtime_task else "unknown"

        # Determine capability type
        capabilities = set(chunk[4] for chunk in chunk_map.get(task_id, []))
        if not capabilities:
            capability_str = "N/A"
        elif len(capabilities) == 1:
            capability_str = list(capabilities)[0]
        else:
            capability_str = "Mixed"

        print(f"🧩 Task: {task_name} (ID: {task_id}, Status: {status}, Capability: {capability_str})")
        
        # Summary count
        chunks_for_task = chunk_map.get(task_id, [])
        total = len(chunks_for_task)
        completed = sum(1 for c in chunks_for_task if c[2] == "completed")
        stale = sum(1 for c in chunks_for_task if c[2] == "stale")
        pending = total - completed - stale
        print(f"   ✓ Chunks: {total} total | {completed} completed | {pending} pending | {stale} stale")

        # Only show chunk info if task is not complete
        if status != "completed":
            for chunk in chunk_map.get(task_id, []):
                chunk_id = chunk[0]
                assigned_agent = chunk[3]
                if assigned_agent:
                    runtime_agent = core._runtime_agents.get(assigned_agent)
                    agent = runtime_agent.obj if runtime_agent else None
                    agent_name = agent.name if agent else assigned_agent[:8]
                else:
                    agent_name = "Unassigned"
                print(f"   └─ Chunk {chunk_id}: {chunk[2]} | Agent: {agent_name}")

        print()  # blank line after each task

def show_zones():
    print("\n🌐 Active Zones:")
    for zone_id, zone in core.zones.items():
        print(f"[Zone {zone.zone_id}] Task {zone.task_id} | Agents: {len(zone.agent_instances)} | Shutdown: {zone.should_shutdown}")

def purge():
    confirm = input("Are you sure you want to purge all data? (yes/no): ")
    if confirm.lower() == 'yes':
        with get_db() as conn:
            conn.execute("DELETE FROM task_logs")
            conn.execute("DELETE FROM task_chunks")
            conn.execute("DELETE FROM tasks")
            conn.execute("DELETE FROM agents")
        from hiveos.runtime.reset_runtime import reset_runtime
        reset_runtime()
        print("💥 System data purged.")
    else:
        print("Purge cancelled.")
        
def fail_agent():
    from hiveos.utilities.induce_failure import induce_failure
    induce_failure(agent_name=None)
    
def demo_progress(demo_prompt):
    input(f"\n{demo_prompt.strip()}\n\n[Press ENTER to continue]\n")
    
def return_injectables():
    print()
    print("'sim-display' ------------- injects simulated display window")
    print("'oled-display' ------------ injects SSD1306 hardware display")
    print("'piezo' ------------------- injects piezo buzzer")
    print("'sim-terminal' ------------ injects simulated terminal window")
    print("'sim-sensor' -------------- injects simulated sensor")
    print()
    print("'agent-shotgun' ----------- injects large batch of randomized sim agents")
    
def set_tick_interval():
    global tick_interval
    try:
        new_interval = float(input("Enter new tick interval in seconds (e.g., 0.2): "))
        if new_interval < 0.01:
            print("Interval too fast. Must be >= 0.01 seconds.")
            return
        tick_interval = new_interval
        print(f"Tick interval set to {tick_interval} seconds.")
    except ValueError:
        print("Invalid input. Please enter a number.")    
    
def inject_yaml_task():
    import os

    filename = input("Enter YAML filename from 'demo_tasks/': ").strip()

    # Automatically append '.yaml' if the user omits it
    if not filename.endswith(".yaml"):
        filename += ".yaml"

    path = os.path.join("hiveos", "demo_tasks", filename)

    if not os.path.exists(path):
        print(f"❌ File not found: {path}")
        return

    with open(path, "r", encoding="utf-8") as f:
        try:
            intent = yaml.safe_load(f)
        except yaml.YAMLError as e:
            print(f"❌ YAML parse error: {e}")
            return

    ingest_intent(intent)
    print(f"✅ Injected YAML task from {filename}")
    
    
def run_demo_mode():
    print("🚀 HiveOS Demo Mode Initiated")

    from hiveos.utilities.induce_failure import induce_failure

    print("\n🧠 Injecting 10 agents...")
    for _ in range(10):
        core.inject_agent(SimAgent(capabilities=["scan"]))
    time.sleep(1)
    
    show_agents()
    print()
    demo_progress("""
        First, we can inject any number of simulated agents with a filterable capability
        These guys have the capability 'scan'
    """)
    
    print()

    print("\n📦 Injecting task with 5 chunks...")
    inject_task(name="DemoTask", count=5, cap="scan")
    time.sleep(1)
    
    show_tasks()
    print()
    demo_progress("""
        Now, we've injected a task that was automatically broken into 5 chunks
        Chunks were given the matching capability of 'scan'
    """)

    print("\n🔄 Ticking system...")
    for _ in range(2):
        tick(verbose=False)
        time.sleep(1)

    demo_progress("""
    After ticking the system, agents have been deployed to a virtual zone controller, spun up
    to handle all local task logic. The zone controller acts as a dynamic edge node, assigning
    agents to appropriate chunks, making requests to the core in the event of shortages, handling
    fallback logic for error states, and shedding excess agents back to the global pool if there   
    is a surplus over a tunable backup agent buffer.
    """)
  
    
    show_agents()
    show_zones()
    
    demo_progress("""
    The idle agent buffer is currently set to 3, and this shows the 3 idle agents tied to the zone controller
    handling DemoTask. They are locked to the zone, and are incapable of being pulled from the global pool until
    released by a zone shed, zone closure, or core priority pull. This ensures that their weighting remains high for
    fallback recovery selection.
    """)

    print("\n💥 Inducing agent failure...")
    induce_failure()
    time.sleep(1)
    
    show_agents()
    
    demo_progress("""
    Now, we can see how failure recovery is handled. By inducing a failure state, like a drained battery in this
    case, an agent enters a "recharging" status, dropping the active chunk to allow for another to pick it up. 
    Any arbitrary failure can trigger this cooldown / handoff, customized by an operator or written in 
    the hardware wrapper. The handoff is seen on the following tick, and work continues with no major interruptions.
    """)
    
    print("\n🔄 Ticking system...")
    tick(verbose=False)
    time.sleep(1)
    
    show_agents()
    
    demo_progress("""
    That same abandoned chunk is now assigned to a different capable agent in the zone's pool. None of this logic
    needed to be sent to the core.
    """)
    
    print("\n📋 Task and Agent Summary:")
    show_tasks()
    show_agents()
    
    demo_progress("""
    Feel free to tick through the rest of this task with selection 3 from the main menu and experiment with 
    the systems. After completion of the DemoTask, the zone that held it will spin down and shed all agents 
    back to the core's global pool. Every event during this demo, and by extension in the entire system, 
    is logged for dumps, metrics, ML datasets, and traceable accountability.
    """)

    
def main():
    verbose_mode = False

    while True:
        print("\nHiveOS CLI")
        print("===========")
        print("1. Inject Task")
        print("2. Inject Agents")
        print("3. Run Tick")
        print("4. Show Agents")
        print("5. Show Tasks")
        print("6. Show Active Zones")
        print("7. Purge All (DB + Runtime)")
        print("8. Toggle Verbose Mode (currently {})".format("ON" if verbose_mode else "OFF"))
        print("9. Induce Random Agent Failure")
        print("10. Show Agent Summaries (Live Introspection)")
        print("11. Remove Agent by ID")
        print("0. Exit")
        print()
        print("'demo' ------------------- run demonstration mode")
        print("'inject-yaml' ------------ inject task from YAML in demo_tasks/")
        print("'s' ---------------------- start continuous tick loop")
        print("'x' ---------------------- stop continuous tick loop")
        print("'more' ------------------- list available special agents")
        print()

        choice = input("Choose an option: ").strip()
        if choice == "1":
            inject_task()
        elif choice == "2":
            inject_agents()
        elif choice == "3":
            run_tick(verbose=verbose_mode)
        elif choice == "4":
            show_agents()
        elif choice == "5":
            show_tasks()
        elif choice == "6":
            show_zones()
        elif choice == "7":
            purge()
        elif choice == "8":
            verbose_mode = not verbose_mode
            print("Verbose mode is now", "ON" if verbose_mode else "OFF")
        elif choice == "9":
            fail_agent()
        elif choice == "10":
            print("\n🧠 Live Agent Summaries")
            for runtime_agent in core._runtime_agents.values():
                agent = runtime_agent.obj
                print("\n-----------------------")
                for key, val in agent.get_summary().items():
                    print(f"{key.title().replace('_', ' ')}: {val}")
        elif choice == "11":
            agent_id = input("Enter full Agent ID to remove: ").strip()
            try:
                if core.remove_agent(agent_id):
                    print(f"Agent '{agent_id}' removed from system.")
                else:
                    print(f"No agent found with ID '{agent_id}'.")
            except Exception as e:
                print(f"❌ Failed to remove agent: {e}") 
        elif choice == "oled-display":
            inject_hw_oled()
        elif choice == "sim-display":
            inject_sim_display() 
        elif choice == "sim-terminal":
            inject_sim_terminal()
        elif choice == "sim-sensor":
            inject_sim_sensor()
        elif choice == "piezo":
            inject_hw_piezo()
        elif choice == "agent-shotgun":
            agent_shotgun()
        elif choice == "demo":
            run_demo_mode()
        elif choice == "more":
            return_injectables()
        elif choice == "inject-yaml":
            inject_yaml_task()
        elif choice == "s":
            global tick_loop_active, tick_thread
            if not tick_loop_active:
                tick_loop_active = True
                tick_thread = threading.Thread(target=background_tick_loop, daemon=True)
                tick_thread.start()
                print("✅ Background tick loop started.")
            else:
                print("⚠️ Loop already running.")            
        elif choice == "x":
            tick_loop_active = False
            print("🛑 Background tick loop stopped.")
        elif choice == "0":
            print("Exiting HiveOS CLI.")
            break
        else:
            print("Invalid option.")

if __name__ == "__main__":
    main()

### USE THIS TO INJECT TERMINAL AGENTS!! ###
# from agents_hardware.terminal_agent import TerminalAgent
#agent = TerminalAgent("AA:BB:CC:DD:EE:FF")  # Replace with real BLE address
