"""
Network Assistant - Let LLM decide what to check
Simple conversational network automation
"""
import json
import ollama
import paramiko

def ask_llm_what_to_check(user_question):
    """Ask LLM what command we should run"""
    prompt = f"""User asked: "{user_question}"
What Cisco IOS command should we run to answer this?
Common commands:
- "show ip interface brief" - check interface status
- "show ip route" - check routing table
- "show version" - check device info
- "show running-config" - check configuration
Return JSON with "command" and "reason" fields."""
    response = ollama.chat(
        model='llama3.2:3b',
        messages=[{'role': 'user', 'content': prompt}],
        format={'type': 'object', 'properties': {'command': {'type': 'string'}, 'reason': {'type': 'string'}}, 'required': ['command', 'reason']},
        options={'temperature': 0}
    )
    return json.loads(response['message']['content'])


def run_command_on_router(host, username, password, command):
    """Execute command on router via SSH"""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(host, username=username, password=password, timeout=10)
    stdin, stdout, stderr = ssh.exec_command(command)
    output = stdout.read().decode()
    ssh.close()
    return output


def ask_llm_to_summarize(command, output):
    """Ask LLM to explain what the output means"""
    prompt = f"""I ran this command: {command}
Output:
{output}
Please summarize what this shows in 2-3 sentences. Focus on any problems or important info."""
    response = ollama.chat(
        model='llama3.2:3b',
        messages=[{'role': 'user', 'content': prompt}],
        options={'temperature': 0}
    )
    return response['message']['content']

def main():
    print("=" * 70)
    print("  NETWORK ASSISTANT")
    print("=" * 70)
    # Load devices from JSON file
    with open('devices.json', 'r') as f:
        devices = json.load(f)
    print(f"\nLoaded {len(devices)} devices: {', '.join(d['name'] for d in devices)}")
    # 1. User question
    user_question = "please verify what is going on with my network"
    print(f"\nYou: {user_question}")
    # 2. Ask LLM what command to run (only once)
    print("\n[LLM is deciding what to check...]")
    decision = ask_llm_what_to_check(user_question)
    print(f"\nLLM decided to run: {decision['command']}")
    print(f"Reason: {decision['reason']}")
    # 3-5. Check each device
    for device in devices:
        print(f"\n{'='*70}")
        print(f"  Checking {device['name']} ({device['host']})")
        print('='*70)
        try:
            # 3. Execute command on router
            output = run_command_on_router(
                device['host'],
                device['username'],
                device['password'],
                decision['command']
            )
            print("\nRouter output:")
            print("-" * 70)
            print(output)
            print("-" * 70)
            # 4. Ask LLM to summarize
            print("\n[LLM is analyzing...]")
            summary = ask_llm_to_summarize(decision['command'], output)
            # 5. Show result
            print(f"\nAssistant: {summary}")
        except Exception as e:
            print(f"\n[ERROR] Could not check {device['name']}: {e}")
    print("\n" + "=" * 70)
    print("  Network check complete")
    print("=" * 70)


if __name__ == "__main__":
    main()