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

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

  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)"; # Defaults to false
      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 = mkMerge [
    # Combined environment packages
    {
      environment.systemPackages =
        (if cfg.enable then [ cfg.package ] else [ ]) ++
        (if cfg.enable && cfg.frontend.enable then [ cfg.frontend.uiPackage cfg.frontend.tikaPackage ] else [ ]);
    }

    # Backend Ollama Service Configuration
    (mkIf cfg.enable {
      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};
            '';
          };
        };
      };
    })

    # Frontend Open WebUI and Tika Service Configuration
    (mkIf (cfg.enable && cfg.frontend.enable) {
      services.open-webui = {
        enable = true;
        package = cfg.frontend.uiPackage;
      };

      services.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;
          };
        };
      };
    })
  ];
}