import platform
import psutil

def hw_detect():
    print("LLM Hardware Check\n", flush=True)

    print(f"OS: {platform.system()} {platform.release()}", flush=True)
    print(f"CPU: {platform.processor()}", flush=True)
    
    try:
        count = psutil.cpu_count(logical=True)
        print(f"Cores: {count}", flush=True)
    except Exception as e:
        print(f"Cores: Error detecting cores ({e})", flush=True)

    try:
        ram = psutil.virtual_memory().total / (1024 ** 3)
        print(f"RAM: {ram:.2f} GB", flush=True)
    except Exception as e:
        print(f"RAM: Error detecting RAM ({e})", flush=True)

    try:
        import torch
        if torch.cuda.is_available():
            try:
                print(f"GPU: {torch.cuda.get_device_name(0)}", flush=True)
            except Exception as e:
                 print(f"GPU: Error getting device name ({e})", flush=True)
        else:
            print("GPU: Not available", flush=True)
    except ImportError:
        print("GPU: torch not installed", flush=True)
    except Exception as e:
        print(f"GPU: Unexpected error ({e})", flush=True)

def recommend():
    print("LLM Model Recommendations\n", flush=True)
    
    try:
        ram_gb = psutil.virtual_memory().total / (1024 ** 3)
        print(f"Detected RAM: {ram_gb:.2f} GB", flush=True)
        
        if ram_gb < 8:
            print("- Recommendation: Use 1-3B models (e.g., Phi-3, Gemma-2B).", flush=True)
        elif ram_gb < 16:
            print("- Recommendation: Use 7-8B models with 4-bit quantization (e.g., Llama-3-8B).", flush=True)
        elif ram_gb < 32:
            print("- Recommendation: Use 14-20B models or 8B models with high precision.", flush=True)
        else:
            print("- Recommendation: You can run 30B+ models or long-context 8-20B models comfortably.", flush=True)
            
        try:
            import torch
            if torch.cuda.is_available():
                vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
                print(f"Detected VRAM: {vram_gb:.2f} GB", flush=True)
                if vram_gb < 6:
                    print("- GPU Note: Limited VRAM. Offload small layers to GPU, keep most in RAM.", flush=True)
                else:
                    print("- GPU Note: Sufficient VRAM for 4-bit 8B models entirely on GPU.", flush=True)
        except ImportError:
            pass
            
    except Exception as e:
        print(f"Error generating recommendations: {e}", flush=True)

def main():
    import sys

    if len(sys.argv) < 2:
        print("Usage: llm-checker [hw-detect | recommend]", flush=True)
        return

    command = sys.argv[1]

    if command == "hw-detect":
        hw_detect()
    elif command == "recommend":
        recommend()
    else:
        print(f"Unknown command: {command}", flush=True)

if __name__ == "__main__":
    main()
