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

let
  inherit (lib) mkIf mkEnableOption mkPackageOption mkOption types optionalAttrs;

  cfg = config.pepe.services.llm;
in
{
  options.pepe.services.llm = {
    enable = mkEnableOption "Enable LLM backend service (Ollama)";

    package = mkPackageOption pkgs "ollama" { };

    backendDomain = mkOption {
      type = types.nullOr types.str;
      default = null;
      description = "Domain for Ollama backend. If null, no vhost is created for the backend.";
    };

    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";
    };

    frontend = {
      enable = mkEnableOption "Enable LLM frontend service (Open WebUI)";
      uiPackage = mkPackageOption pkgs "open-webui" { };
      tikaPackage = mkPackageOption pkgs "tika" { }; # Tika for document processing with Open WebUI
      domain = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Domain for Open WebUI frontend. If null, no vhost is created for the frontend.";
      };
    };
  };

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

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

      pepe.core.vhost.hosts = optionalAttrs (cfg.backendDomain != null) {
        "${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; # Ollama can take time to respond
              proxy_set_header Host localhost:${toString config.services.ollama.port};
            '';
          };
        };
      };
    }
  // mkIf (cfg.enable && cfg.frontend.enable) {
    environment.systemPackages = [
      cfg.frontend.uiPackage
      cfg.frontend.tikaPackage
    ];

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

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

    pepe.core.vhost.hosts = optionalAttrs (cfg.frontend.domain != null) {
      "${cfg.frontend.domain}" = {
        locations."/" = {
          host = config.services.open-webui.host;
          port = config.services.open-webui.port;
          allowLAN = true;
          allowVPN = true;
          allowWAN = true;
          proxyWebsockets = true;
        };
      };
    };
  };
}