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

let
  inherit (lib) mkIf;

  cfg = config.pepe.services.llm;
in
{
  options.pepe.services.llm = with lib; {
    enable = mkEnableOption "Enable LLM services (Ollama and Open WebUI)";
    package = mkPackageOption pkgs "ollama" { };
    uiPackage = mkPackageOption pkgs "open-webui" { };
    tikaPackage = mkPackageOption pkgs "tika" { };
    backendDomain = mkOption {
      type = types.str;
      default = null;
    };
    frontendDomain = mkOption {
      type = types.str;
      default = null;
    };
    acceleration = mkOption {
      type = types.enum [ "none" "cuda" "rocm" ];
      default = "none";
      description = "Acceleration type for Ollama";
    };
    environmentVariables = mkOption {
      type = types.attrsOf types.str;
      default = {
        OLLAMA_FLASH_ATTENTION = "1";
        OLLAMA_NUM_PARALLEL = "2";
        OLLAMA_KV_CACHE_TYPE = "q8_0";
      };
      description = "Environment variables for Ollama";
    };
  };

  config = mkIf cfg.enable {
    environment.systemPackages = [ cfg.package ];

    services = {
      ollama = {
        enable = true;
        package = cfg.package;
        acceleration = cfg.acceleration;
        environmentVariables = cfg.environmentVariables;
      };

      open-webui = {
        enable = true;
        package = cfg.uiPackage;
      };

      tika = {
        enable = true;
        package = cfg.tikaPackage;
      };
    };

    pepe.core.vhost.hosts.${cfg.backendDomain} = {
      locations."/" = {
        host = config.services.ollama.host;
        port = config.services.ollama.port;
        allowLAN = true;
        allowVPN = true;
        allowWAN = true;
        recommendedProxySettings = false;
        extraConfig = ''
          proxy_buffering off;
          proxy_read_timeout 600s;
          proxy_set_header Host localhost:${toString config.services.ollama.port};
        '';
      };
    };

    pepe.core.vhost.hosts.${cfg.frontendDomain} = {
      locations."/" = {
        host = config.services.open-webui.host;
        port = config.services.open-webui.port;
        allowLAN = true;
        allowVPN = true;
        allowWAN = true;
        proxyWebsockets = true;
      };
    };
  };
}