{ config, pkgs, lib, ... }:

let
  # Read environment variables (OLLAMA_HOST and MISTRAL_API_KEY) from a local file.
  # This keeps sensitive values out of the Nix store and version control.
  # File format: one KEY=value per line (e.g., MISTRAL_API_KEY=yourkey)
  envVars = lib.genAttrs (builtins.splitStrings "\n" (builtins.readFile (toString ./assets/conf/apps/ai.env)))
    (nameValue: builtins.splitString "=" nameValue);
in
{
  # --- Packages ---
  # Install ZED editor and Ollama with Vulkan support (for CPU/AMD).
  # For NVIDIA GPUs, replace `ollama-vulkan` with `ollama`.
  # For AMD ROCm, use `ollama-rocm` and ensure ROCm is installed.
  home.packages = [
    pkgs.ollama-vulkan
    pkgs.zed
  ];

  # --- Environment Variables ---
  # Set OLLAMA_HOST and MISTRAL_API_KEY for ZED and other user applications.
  # Values are read from ./assets/conf/apps/ai.env.
  home.sessionVariables = {
    OLLAMA_HOST = envVars.OLLAMA_HOST or "http://127.0.0.1:11434";  # Default Ollama endpoint
    MISTRAL_API_KEY = envVars.MISTRAL_API_KEY or "";               # Mistral API key (required for cloud models)
  };

  # --- Ollama User Service ---
  # Configure Ollama to run as a user service (starts on login).
  # This avoids root privileges and allows per-user model management.
  systemd.user.services.ollama = {
    description = "Ollama service for local AI models";
    wantedBy = [ "default.target" ];  # Start with user session
    after = [ "network.target" ];     # Ensure network is ready

    serviceConfig = {
      Type = "forking";  # Run as a background process

      # Start Ollama server
      ExecStart = ''
        ${pkgs.ollama-vulkan}/bin/ollama serve
      '';

      # Pull default models after server starts
      ExecStartPost = ''
        sleep 5  # Wait for server to initialize

        # Pull coding and chat models at startup
        ${pkgs.ollama-vulkan}/bin/ollama pull codellama:70b    # Best for coding tasks
        ${pkgs.ollama-vulkan}/bin/ollama pull mixtral:8x7b    # Best for chat/conversation

        # Uncomment to pull additional models:
        # ${pkgs.ollama-vulkan}/bin/ollama pull llama3:8b       # General-purpose model
        # ${pkgs.ollama-vulkan}/bin/ollama pull qwen2.5-coder:7b # Multilingual coding
        # ${pkgs.ollama-vulkan}/bin/ollama pull starcoder2:15b  # Alternative for code
      '';
      Restart = "on-failure";  # Restart if Ollama crashes
    };
  };

  # --- ZED Configuration ---
  # Generate ZED's settings.json with substituted API keys and endpoints.
  # Base config is read from ./assets/conf/apps/ai.conf, with variables injected.
  home.file.".config/zed/settings.json".text = lib.mkForce (
    builtins.readFile (toString ./assets/conf/apps/ai.conf)
    // ''
      {
        "mistral": {
          "apiKey": "${envVars.MISTRAL_API_KEY}",  # Inject Mistral API key
          "defaultModel": "mistral-pro"            # Default Mistral model
        },
        "ollama": {
          "endpoint": "${envVars.OLLAMA_HOST}",    # Inject Ollama endpoint
          "defaultModel": "codellama:70b"          # Default Ollama model for coding
        }
      }
    ''
  );

  # --- Usage Notes ---
  # 1. Pulling Additional Models:
  #    To add more models later, run:
  #      ollama pull <model-name>
  #    Example: ollama pull llama3:8b

  # 2. Switching GPU Backends:
  #    - For NVIDIA: Replace all `ollama-vulkan` with `ollama` (uses CUDA)
  #    - For AMD: Use `ollama-rocm` and ensure ROCm is installed

  # 3. ZED Plugin Setup:
  #    - Install the Ollama and Mistral plugins in ZED via the plugin marketplace
  #    - The Ollama plugin will use the local models pulled above
  #    - The Mistral plugin will use the MISTRAL_API_KEY for cloud access

  # 4. Security:
  #    - Never commit ./assets/conf/apps/ai.env to version control
  #    - For extra security, encrypt ai.env using sops-nix or age

  # 5. Persistent Service:
  #    To keep Ollama running after logout, enable lingering:
  #      loginctl enable-linger $(whoami)
}