#!/usr/bin/env python3
import sys
import shutil
import subprocess
import torch
import os

def check_flatpak():
    return os.path.exists("/.flatpak-info")

def run_cmd(cmd_list):
    in_flatpak = check_flatpak()
    if in_flatpak:
        full_cmd = ["flatpak-spawn", "--host"] + cmd_list
    else:
        full_cmd = cmd_list
    
    try:
        result = subprocess.run(full_cmd, capture_output=True, text=True)
        return result.returncode == 0, result.stdout.strip(), result.stderr.strip()
    except FileNotFoundError:
        return False, "", "Command not found"
    except Exception as e:
        return False, "", str(e)

def main():
    print("========================================")
    print("   AI Transcriber Diagnostic Tool")
    print("========================================")
    
    in_flatpak = check_flatpak()
    print(f"Environment: {'Flatpak Sandbox' if in_flatpak else 'Native Host'}")
    print("-" * 40)

    # 1. GPU Check (PyTorch)
    print("\n[1] GPU Availability (Internal PyTorch)")
    if torch.cuda.is_available():
        print(f"✅ GPU Detected: {torch.cuda.get_device_name(0)}")
        print(f"   VRAM: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
        print(f"   CUDA Version: {torch.version.cuda}")
    else:
        print("❌ GPU NOT Detected by PyTorch.")
        print("   Whisper will run on CPU (Slow).")

    # 2. Nvidia Driver Check (System)
    print("\n[2] Nvidia Driver Check (System)")
    ok, out, err = run_cmd(["nvidia-smi", "--query-gpu=name,driver_version", "--format=csv,noheader"])
    if ok:
        print(f"✅ Driver Active: {out}")
    else:
        print("❌ 'nvidia-smi' failed. Drivers might be missing or inaccessible.")
        if err: print(f"   Error: {err}")

    # 3. FFmpeg Check
    print("\n[3] FFmpeg Check")
    ok, out, err = run_cmd(["ffmpeg", "-version"])
    if ok:
        version_line = out.split('\n')[0]
        print(f"✅ FFmpeg Ready: {version_line}")
    else:
        print("❌ FFmpeg NOT found.")
        print("   On Bazzite, install it with: 'brew install ffmpeg'")

    # 4. FFprobe Check
    print("\n[4] FFprobe Check (Required for duration checks)")
    ok, out, err = run_cmd(["ffprobe", "-version"])
    if ok:
        version_line = out.split('\n')[0]
        print(f"✅ FFprobe Ready: {version_line}")
    else:
        print("❌ FFprobe NOT found.")

    # 5. Ollama Check
    print("\n[5] Local LLM (Ollama)")
    ok, out, err = run_cmd(["ollama", "--version"])
    if ok:
        print(f"✅ Ollama Installed: {out}")
        
        # Check server
        import socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(1)
        result = sock.connect_ex(('127.0.0.1', 11434))
        if result == 0:
            print("✅ Ollama Server is RUNNING.")
        else:
            print("⚠️  Ollama Server is NOT running (Scripts will try to auto-start it).")
    else:
        print("❌ Ollama NOT found.")
        print("   Run: ./ai_transcriber_v2/install_local_llm.sh")

    print("\n========================================")

if __name__ == "__main__":
    main()