require "json"
require "./data_files"

module JosieHealth::LUT
  module CaffeineLUT
    # Serving sizes in ml
    SERVING_SIZES = {
      "can"     => 355,
      "cans"    => 355,
      "bottle"  => 500,
      "bottles" => 500,
      "cup"     => 240,
      "cups"    => 240,
      "mug"     => 350,
      "mugs"    => 350,
      "shot"    => 30,
      "shots"   => 30,
    }

    # Category defaults
    CATEGORY_DEFAULTS = {
      "energy_drink" => {caffeine_mg: 160, default_ml: 500},
      "energy_shot"  => {caffeine_mg: 200, default_ml: 60},
      "coffee"       => {caffeine_mg: 95, default_ml: 240},
      "tea"          => {caffeine_mg: 50, default_ml: 240},
      "soda"         => {caffeine_mg: 40, default_ml: 355},
    }

    @@brands : Hash(String, NamedTuple(caffeine_mg: Int32, category: String, default_ml: Int32))?
    @@powder_data : Hash(String, Int32)?
    @@alias_map : Hash(String, String)?
    @@loaded = false

    # Load caffeine data from embedded JSON
    def self.load_data
      return if @@loaded

      begin
        json = DataFiles.caffeine

        @@brands = {} of String => NamedTuple(caffeine_mg: Int32, category: String, default_ml: Int32)
        @@powder_data = {} of String => Int32
        @@alias_map = {} of String => String

        if brands_json = json["brands"]?
          brands_json.as_h.each do |name, data|
            canonical_name = name.to_s
            @@brands.not_nil![canonical_name] = {
              caffeine_mg: data["caffeine_mg"].as_i,
              category:    data["category"].as_s,
              default_ml:  data["default_ml"].as_i,
            }

            if mg_per_g = data["caffeine_mg_per_g"]?
              @@powder_data.not_nil![canonical_name] = mg_per_g.as_i
            end

            # Register aliases
            if aliases = data["aliases"]?
              aliases.as_a.each do |alias_name|
                @@alias_map.not_nil![alias_name.as_s] = canonical_name
              end
            end
          end
        end

        puts "Loaded #{@@brands.not_nil!.size} caffeine brands with #{@@alias_map.not_nil!.size} aliases"
        @@loaded = true
      rescue ex
        puts "Error loading caffeine.json: #{ex.message}"
        @@brands = {} of String => NamedTuple(caffeine_mg: Int32, category: String, default_ml: Int32)
        @@powder_data = {} of String => Int32
        @@alias_map = {} of String => String
        @@loaded = true
      end
    end

    # Resolve alias to canonical name
    def self.resolve_alias(name : String) : String
      load_data
      normalized = name.downcase.gsub(/[\s\-_]/, "")
      @@alias_map.not_nil![normalized]? || normalized
    end

    # Lookup a brand by name (handles aliases)
    def self.lookup_brand(name : String) : NamedTuple(caffeine_mg: Int32, category: String, default_ml: Int32)?
      load_data
      canonical = resolve_alias(name)
      @@brands.not_nil![canonical]?
    end

    # Check if a substance is a caffeinated beverage/product
    def self.is_caffeine_source?(name : String) : Bool
      lookup_brand(name) != nil
    end

    # Lookup caffeine mg per gram for powder sources (e.g., matcha)
    def self.lookup_powder_mg_per_g(name : String) : Int32?
      load_data
      canonical = resolve_alias(name)
      @@powder_data.not_nil![canonical]?
    end

    # Get category defaults
    def self.get_category_defaults(category : String) : NamedTuple(caffeine_mg: Int32, default_ml: Int32)?
      CATEGORY_DEFAULTS[category.downcase]?
    end

    # Calculate caffeine mg from volume and brand info
    def self.calculate_mg(volume_ml : Float64, caffeine_per_serving : Int32, serving_ml : Int32) : Float64
      ((volume_ml / serving_ml) * caffeine_per_serving).round(1)
    end

    # Parse dosage string like "2cans", "500ml", "1cup"
    # Returns {count, serving_type, volume_ml} or nil if not parseable
    def self.parse_dosage(dosage : String) : NamedTuple(count: Int32, serving: String?, volume_ml: Int32)?
      normalized = dosage.downcase.strip

      # Try to match "Nunit" pattern (e.g., "2cans", "1cup", "500ml")
      if match = normalized.match(/^(\d+(?:\.\d+)?)(ml|can|cans|bottle|bottles|cup|cups|mug|mugs|shot|shots)?$/)
        count_str = match[1]
        serving = match[2]?

        # Handle decimal counts
        count_f = count_str.to_f
        count = count_f.to_i

        if serving == "ml"
          # Direct ml specification
          return {count: 1, serving: "ml", volume_ml: count_f.to_i}
        elsif serving && SERVING_SIZES.has_key?(serving)
          volume_ml = SERVING_SIZES[serving] * count
          return {count: count, serving: serving, volume_ml: volume_ml}
        elsif serving.nil?
          # Just a number, assume count of default serving
          return {count: count, serving: nil, volume_ml: 0}
        end
      end

      nil
    end

    # Full caffeine dose calculation
    # Input: dosage string, substance name
    # Returns: {caffeine_mg, annotation} or nil
    def self.calculate_dose(dosage : String, substance : String) : NamedTuple(caffeine_mg: Float64, annotation: String)?
      parsed = parse_dosage(dosage)
      return nil unless parsed

      brand_info = lookup_brand(substance)
      return nil unless brand_info

      # Determine volume
      volume_ml = if parsed[:volume_ml] > 0
                    parsed[:volume_ml]
                  else
                    parsed[:count] * brand_info[:default_ml]
                  end

      caffeine_mg = calculate_mg(volume_ml.to_f, brand_info[:caffeine_mg], brand_info[:default_ml])

      # Build details string
      serving_str = parsed[:serving] ? "#{parsed[:count]}#{parsed[:serving]}" : parsed[:count].to_s
      canonical = resolve_alias(substance)
      details = "#{canonical} #{serving_str} (#{volume_ml}ml) = #{caffeine_mg.to_i}mg"

      {caffeine_mg: caffeine_mg, annotation: details}
    end

    # Get all brand names
    def self.all_brands : Array(String)
      load_data
      @@brands.not_nil!.keys
    end
  end
end