require 'mini_magick'
require_relative 'substance_categorizer'

class GraphGenerator
  def initialize(width = 800, height = 400)
    @width = width
    @height = height + 100  # Add extra height for legend space
    @margin = 60
    @title_margin = 50  # Space between title and chart
    @legend_height = 80  # Space reserved for legend
    @chart_width = @width - (@margin * 2)
    @chart_height = height - (@margin * 2) - @legend_height - @title_margin  # Reduce chart area for title and legend
  end

  def generate_monthly_doses_line_graph(monthly_data)
    # Use OUTPUT_DIR environment variable or current directory for temporary files
    output_dir = ENV['OUTPUT_DIR'] || '.'
    temp_file = File.join(output_dir, 'temp_line_graph_base.png')

    # Create high-quality image using MiniMagick::Tool
    MiniMagick::Tool.new('magick') do |convert|
      convert.size "#{@width}x#{@height}"
      convert.density "150"
      convert.depth "8"
      convert.colorspace "sRGB"
      convert << 'xc:transparent'
      convert << temp_file
    end

    image = MiniMagick::Image.open(temp_file)

    image.combine_options do |c|
      # Draw axes with thicker lines
      c.stroke 'white'
      c.strokewidth 3
      c.fill 'none'
      
      chart_top = @margin + @title_margin
      chart_bottom = chart_top + @chart_height
      
      # Y-axis
      c.draw "line #{@margin},#{chart_top} #{@margin},#{chart_bottom}"
      # X-axis
      c.draw "line #{@margin},#{chart_bottom} #{@width - @margin},#{chart_bottom}"
      
      # Prepare data points
      months = monthly_data.keys.sort
      return image if months.empty?
      
      max_doses = monthly_data.values.map { |data| data.values.sum { |substance_data| substance_data[:count] } }.max
      return image if max_doses == 0
      
      points = []
      months.each_with_index do |month, index|
        # Fix division by zero when months.length <= 1
        x = if months.length == 1
              @margin + @chart_width / 2  # Center single point
            else
              @margin + (index.to_f / (months.length - 1)) * @chart_width
            end
        total_doses = monthly_data[month].values.sum { |substance_data| substance_data[:count] }
        y = chart_bottom - (total_doses.to_f / max_doses) * @chart_height
        points << [x, y]
      end
      
      # Draw line graph with Spotify colors
      if points.length > 1
        c.stroke 'rgb(29,185,84)'  # Spotify green
        c.strokewidth 5
        c.fill 'none'
        
        line_path = "M #{points[0][0]},#{points[0][1]}"
        points[1..-1].each do |point|
          line_path += " L #{point[0]},#{point[1]}"
        end
        c.draw "path '#{line_path}'"
        
        # Draw points with Spotify green
        points.each do |point|
          c.stroke 'none'
          c.fill 'rgb(29,185,84)'
          c.draw "circle #{point[0]},#{point[1]} #{point[0] + 6},#{point[1]}"
        end
      end
      
      # Add labels - no stroke, just fill
      c.stroke 'none'
      c.pointsize 24
      c.fill 'white'
      c.font 'DejaVu-Sans'
      c.gravity 'Center'
      
      # Month labels on X-axis
      c.gravity 'North'
      months.each_with_index do |month, index|
        # Fix division by zero when months.length <= 1
        x = if months.length == 1
              @margin + @chart_width / 2  # Center single point
            else
              @margin + (index.to_f / (months.length - 1)) * @chart_width
            end
        month_name = Date.strptime(month, '%Y-%m').strftime('%b')
        # Fix NaN issue in annotate coordinates
        x_offset = (x - @width/2).to_i
        c.annotate "+#{x_offset}+#{chart_bottom + 10}", month_name
      end
      
      # Y-axis labels
      c.gravity 'NorthWest'
      c.stroke 'none'
      c.fill 'white'
      c.pointsize 20
      c.font 'DejaVu-Sans'
      (0..4).each do |i|
        value = (max_doses * i / 4).round
        y = chart_bottom - (i.to_f / 4) * @chart_height
        c.annotate "+#{5}+#{y - 10}", value.to_s
      end
      
      # Title - positioned at top left with more margin
      c.stroke 'none'
      c.pointsize 32
      c.fill 'white'
      c.font 'DejaVu-Sans-Bold'
      c.gravity 'NorthWest'
      c.annotate "+#{@margin}+20", "Monthly Dose Count"
    end
    
    # Clean up temporary file
    # Clean up temporary file
    output_dir = ENV['OUTPUT_DIR'] || '.'
    temp_file = File.join(output_dir, 'temp_line_graph_base.png')
    File.delete(temp_file) if File.exist?(temp_file)
    
    image
  end

  def generate_category_stacked_chart(monthly_category_data)
    # Create high-quality image using MiniMagick::Tool
    MiniMagick::Tool.new('magick') do |convert|
      convert.size "#{@width}x#{@height}"
      convert.density "150"
      convert.depth "8"
      convert.colorspace "sRGB"
      convert << 'xc:transparent'
      convert << 'temp_category_chart_base.png'
    end
    
    image = MiniMagick::Image.open('temp_category_chart_base.png')

    image.combine_options do |c|
      months = monthly_category_data.keys.sort
      return image if months.empty?
      
      categories = SubstanceCategorizer.get_categories
      
      # Spotify-inspired color palette with fixed mapping per category
      category_colors = {
        "Psychedelics" => 'rgb(240,55,165)',     # Spotify pink
        "Stimulants" => 'rgb(255,167,38)',       # Spotify orange  
        "Depressants" => 'rgb(66,165,245)',      # Spotify blue
        "Dissociatives" => 'rgb(171,71,188)',    # Spotify purple
        "Opioids" => 'rgb(239,83,80)',           # Spotify red
        "Cannabis" => 'rgb(29,185,84)',          # Spotify green
        "Nootropics" => 'rgb(38,198,218)',       # Spotify cyan
        "Medications" => 'rgb(120,144,156)',     # Spotify grey
        "Supplements" => 'rgb(102,187,106)',     # Light green
        "Botanicals" => 'rgb(30,215,96)',        # Lighter Spotify green
        "Inhalants" => 'rgb(255,202,40)',        # Spotify yellow
        "Research Chemicals" => 'rgb(141,110,99)', # Spotify brown
        "Other" => 'rgb(200,200,200)'            # Light grey
      }
      
      # Calculate maximum total for scaling
      max_total = 0
      months.each do |month|
        total = categories.sum { |cat| monthly_category_data[month][cat] || 0 }
        max_total = [max_total, total].max
      end
      
      return image if max_total == 0
      
      # Draw axes
      c.stroke 'white'
      c.strokewidth 3
      c.fill 'none'
      chart_top = @margin + @title_margin
      chart_bottom = chart_top + @chart_height
      c.draw "line #{@margin},#{chart_top} #{@margin},#{chart_bottom}"
      c.draw "line #{@margin},#{chart_bottom} #{@width - @margin},#{chart_bottom}"
      
      # Draw stacked bars
      bar_width = @chart_width / months.length * 0.8
      months.each_with_index do |month, month_index|
        x = @margin + (month_index + 0.1) * (@chart_width / months.length)
        current_y = chart_bottom
        
        categories.each_with_index do |category, cat_index|
          count = monthly_category_data[month][category] || 0
          next if count == 0
          
          segment_height = (count.to_f / max_total) * @chart_height
          c.stroke 'none'
          c.fill category_colors[category] || 'rgb(200,200,200)'
          c.draw "rectangle #{x},#{current_y - segment_height} #{x + bar_width},#{current_y}"
          current_y -= segment_height
        end
      end
      
      # Month labels - no stroke, just fill
      c.stroke 'none'
      c.pointsize 24
      c.fill 'white'
      c.font 'DejaVu-Sans'
      c.gravity 'Center'
      
      c.gravity 'North'
      months.each_with_index do |month, index|
        x = @margin + (index + 0.5) * (@chart_width / months.length)
        month_name = Date.strptime(month, '%Y-%m').strftime('%b')
        c.annotate "+#{x - @width/2}+#{chart_bottom + 10}", month_name
      end
      
      # Y-axis labels
      c.gravity 'NorthWest'
      c.stroke 'none'
      c.fill 'white'
      c.pointsize 20
      c.font 'DejaVu-Sans'
      (0..4).each do |i|
        value = (max_total * i / 4).round
        y = chart_bottom - (i.to_f / 4) * @chart_height
        c.annotate "+#{5}+#{y - 10}", value.to_s
      end
      
      # Title - positioned at top left with more margin
      c.stroke 'none'
      c.pointsize 32
      c.fill 'white'
      c.font 'DejaVu-Sans-Bold'
      c.gravity 'NorthWest'
      c.annotate "+#{@margin}+20", "Categories by Month"
      
      # Legend positioned under the graph with proper boundaries
      legend_start_y = chart_bottom + 50
      
      # Find categories that actually have data
      categories_with_data = []
      months.each do |month|
        categories.each do |category|
          count = monthly_category_data[month][category] || 0
          if count > 0 && !categories_with_data.include?(category)
            categories_with_data << category
          end
        end
      end
      
      items_per_row = 2
      legend_item_width = (@chart_width - 40) / items_per_row  # Ensure fit within chart width
      
      categories_with_data.each_with_index do |category, index|
        row = index / items_per_row
        col = index % items_per_row
        legend_x = @margin + 20 + col * legend_item_width
        legend_y = legend_start_y + row * 30
        
        # Draw colored rectangle
        c.stroke 'none'
        c.fill category_colors[category] || 'rgb(200,200,200)'
        c.draw "rectangle #{legend_x},#{legend_y} #{legend_x + 20},#{legend_y + 20}"
        
        # Legend text
        c.stroke 'none'
        c.fill 'white'
        c.pointsize 18
        c.font 'DejaVu-Sans'
        c.gravity 'NorthWest'
        c.annotate "+#{legend_x + 30}+#{legend_y + 5}", category
      end
    end
    
    # Clean up temporary file
    File.delete('temp_category_chart_base.png') if File.exist?('temp_category_chart_base.png')
    
    image
  end

  def calculate_monthly_category_data(data)
    monthly_category_data = {}
    
    # Group data by month and category
    data.each do |entry|
      month_key = entry[:timestamp].strftime('%Y-%m')
      category = SubstanceCategorizer.categorize_substance(entry[:substance])
      
      monthly_category_data[month_key] ||= {}
      monthly_category_data[month_key][category] ||= 0
      monthly_category_data[month_key][category] += 1
    end
    
    monthly_category_data
  end
end
