package main

import (
	"bytes"
	"database/sql"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/google/uuid"
	"github.com/jhillyerd/enmime"
)

type Config struct {
	EmailDir      string
	VaultDir      string
	ProcessedDir  string
	DBHost        string
	DBUser        string
	DBPassword    string
	DBName        string
}

type EmailProcessor struct {
	db           *sql.DB
	config       Config
	enquiryMap   map[string]int
	invoiceMap   map[string]int
	poMap        map[string]int
	userMap      map[string]int
	jobMap       map[string]int
}

type Attachment struct {
	Type          string
	Name          string
	Filename      string
	Size          int64
	IsMessageBody int
}

func main() {
	var config Config
	flag.StringVar(&config.EmailDir, "emaildir", "/var/www/emails", "Email storage directory")
	flag.StringVar(&config.VaultDir, "vaultdir", "/var/www/vaultmsgs/new", "Vault messages directory")
	flag.StringVar(&config.ProcessedDir, "processeddir", "/var/www/vaultmsgs/cur", "Processed messages directory")
	flag.StringVar(&config.DBHost, "dbhost", "127.0.0.1", "Database host")
	flag.StringVar(&config.DBUser, "dbuser", "cmc", "Database user")
	flag.StringVar(&config.DBPassword, "dbpass", "xVRQI&cA?7AU=hqJ!%au", "Database password")
	flag.StringVar(&config.DBName, "dbname", "cmc", "Database name")
	flag.Parse()

	dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?parseTime=true",
		config.DBUser, config.DBPassword, config.DBHost, config.DBName)
	
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Fatal("Failed to connect to database:", err)
	}
	defer db.Close()

	processor := &EmailProcessor{
		db:     db,
		config: config,
	}

	if err := processor.loadMaps(); err != nil {
		log.Fatal("Failed to load maps:", err)
	}

	if err := processor.processEmails(); err != nil {
		log.Fatal("Failed to process emails:", err)
	}
}

func (p *EmailProcessor) loadMaps() error {
	p.enquiryMap = make(map[string]int)
	p.invoiceMap = make(map[string]int)
	p.poMap = make(map[string]int)
	p.userMap = make(map[string]int)
	p.jobMap = make(map[string]int)

	// Load enquiries
	rows, err := p.db.Query("SELECT id, title FROM enquiries")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var id int
		var title string
		if err := rows.Scan(&id, &title); err != nil {
			return err
		}
		p.enquiryMap[title] = id
	}

	// Load invoices
	rows, err = p.db.Query("SELECT id, title FROM invoices")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var id int
		var title string
		if err := rows.Scan(&id, &title); err != nil {
			return err
		}
		p.invoiceMap[title] = id
	}

	// Load purchase orders
	rows, err = p.db.Query("SELECT id, title FROM purchase_orders")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var id int
		var title string
		if err := rows.Scan(&id, &title); err != nil {
			return err
		}
		p.poMap[title] = id
	}

	// Load users
	rows, err = p.db.Query("SELECT id, email FROM users")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var id int
		var email string
		if err := rows.Scan(&id, &email); err != nil {
			return err
		}
		p.userMap[strings.ToLower(email)] = id
	}

	// Load jobs
	rows, err = p.db.Query("SELECT id, title FROM jobs")
	if err != nil {
		return err
	}
	defer rows.Close()
	for rows.Next() {
		var id int
		var title string
		if err := rows.Scan(&id, &title); err != nil {
			return err
		}
		p.jobMap[title] = id
	}

	return nil
}

func (p *EmailProcessor) processEmails() error {
	files, err := ioutil.ReadDir(p.config.VaultDir)
	if err != nil {
		return err
	}

	for _, file := range files {
		if file.IsDir() || file.Name() == "." || file.Name() == ".." {
			continue
		}

		fmt.Printf("Handling %s\n", file.Name())
		if err := p.processEmail(file.Name()); err != nil {
			log.Printf("Error processing %s: %v\n", file.Name(), err)
		}
	}

	return nil
}

func (p *EmailProcessor) processEmail(filename string) error {
	emailPath := filepath.Join(p.config.VaultDir, filename)
	content, err := ioutil.ReadFile(emailPath)
	if err != nil {
		return err
	}

	if len(content) == 0 {
		fmt.Println("No content found. Ignoring this email")
		return p.moveEmail(filename)
	}

	// Parse email with enmime
	env, err := enmime.ReadEnvelope(bytes.NewReader(content))
	if err != nil {
		return err
	}

	// Get recipients
	toRecipients := p.parseEnmimeAddresses(env.GetHeader("To"))
	fromRecipients := p.parseEnmimeAddresses(env.GetHeader("From"))
	ccRecipients := p.parseEnmimeAddresses(env.GetHeader("Cc"))

	// Check if we should save this email
	saveThis := false
	fromKnownUser := false

	for _, email := range toRecipients {
		if p.userExists(email) {
			saveThis = true
		}
	}

	for _, email := range fromRecipients {
		if p.userExists(email) {
			saveThis = true
			fromKnownUser = true
		}
	}

	for _, email := range ccRecipients {
		if p.userExists(email) {
			saveThis = true
		}
	}

	subject := env.GetHeader("Subject")
	if subject == "" {
		fmt.Println("No subject found. Ignoring this email")
		return p.moveEmail(filename)
	}

	// Check for identifiers in subject
	foundEnquiries := p.checkIdentifier(subject, p.enquiryMap, "enquiry")
	foundInvoices := p.checkIdentifier(subject, p.invoiceMap, "invoice")
	foundPOs := p.checkIdentifier(subject, p.poMap, "purchaseorder")
	foundJobs := p.checkIdentifier(subject, p.jobMap, "job")

	foundIdent := len(foundEnquiries) > 0 || len(foundInvoices) > 0 || 
		len(foundPOs) > 0 || len(foundJobs) > 0

	if fromKnownUser || saveThis || foundIdent {
		// Process and save the email
		unixTime := time.Now().Unix()
		if date, err := env.Date(); err == nil {
			unixTime = date.Unix()
		}

		// Get recipient user IDs
		recipientIDs := make(map[string][]int)
		recipientIDs["to"] = p.getUserIDs(toRecipients)
		recipientIDs["from"] = p.getUserIDs(fromRecipients)
		recipientIDs["cc"] = p.getUserIDs(ccRecipients)

		if len(recipientIDs["from"]) == 0 {
			fmt.Println("Email has no From Recipient ID. Ignoring this email")
			return p.moveEmail(filename)
		}

		fmt.Println("---------START MESSAGE -----------------")
		fmt.Printf("Subject: %s\n", subject)

		// Extract attachments using enmime
		relativePath := p.getAttachmentDirectory(unixTime)
		attachments := p.extractEnmimeAttachments(env, relativePath)

		// Save email to database
		if err := p.saveEmail(filename, subject, unixTime, recipientIDs, attachments,
			foundEnquiries, foundInvoices, foundPOs, foundJobs); err != nil {
			return err
		}

		fmt.Println("--------END MESSAGE ------")
	} else {
		fmt.Printf("Email will not be saved. Subject: %s\n", subject)
	}

	return p.moveEmail(filename)
}

func (p *EmailProcessor) parseEnmimeAddresses(header string) []string {
	var emails []string
	if header == "" {
		return emails
	}

	addresses, err := enmime.ParseAddressList(header)
	if err != nil {
		return emails
	}

	for _, addr := range addresses {
		emails = append(emails, strings.ToLower(addr.Address))
	}

	return emails
}

func (p *EmailProcessor) userExists(email string) bool {
	_, exists := p.userMap[strings.ToLower(email)]
	return exists
}

func (p *EmailProcessor) getUserIDs(emails []string) []int {
	var ids []int
	for _, email := range emails {
		if id, exists := p.userMap[strings.ToLower(email)]; exists {
			ids = append(ids, id)
		} else {
			// Create new user
			newID := p.createUser(email)
			if newID > 0 {
				ids = append(ids, newID)
				p.userMap[strings.ToLower(email)] = newID
			}
		}
	}
	return ids
}

func (p *EmailProcessor) createUser(email string) int {
	fmt.Printf("Making a new User for: '%s'\n", email)
	
	result, err := p.db.Exec(
		"INSERT INTO users (type, email, by_vault, created, modified) VALUES (?, ?, ?, NOW(), NOW())",
		"contact", strings.ToLower(email), 1)
	
	if err != nil {
		fmt.Printf("Serious Error: Unable to create user for email '%s': %v\n", email, err)
		return 0
	}

	id, err := result.LastInsertId()
	if err != nil {
		return 0
	}

	fmt.Printf("New User '%s' Added with ID: %d\n", email, id)
	return int(id)
}

func (p *EmailProcessor) checkIdentifier(subject string, identMap map[string]int, identType string) []int {
	var results []int
	var re *regexp.Regexp

	switch identType {
	case "enquiry":
		re = regexp.MustCompile(`CMC\d+([NVQWSOT]|ACT|NT)E\d+-\d+`)
	case "invoice":
		re = regexp.MustCompile(`CMCIN\d+`)
	case "purchaseorder":
		re = regexp.MustCompile(`CMCPO\d+`)
	case "job":
		re = regexp.MustCompile(`(JAN|FEB|MAR|APR|MAY|JUN|JUL|AUG|SEP|OCT|NOV|DEC)\d+(N|V|W|S|T|NT|ACT|Q|O)J\d+`)
	}

	if re != nil {
		matches := re.FindAllString(subject, -1)
		for _, match := range matches {
			if id, exists := identMap[match]; exists {
				results = append(results, id)
			}
		}
	}

	return results
}

func (p *EmailProcessor) getAttachmentDirectory(unixTime int64) string {
	t := time.Unix(unixTime, 0)
	monthYear := t.Format("01-2006")
	path := filepath.Join(p.config.EmailDir, monthYear)

	if err := os.MkdirAll(path, 0755); err != nil {
		log.Printf("Failed to create directory %s: %v", path, err)
	}

	return monthYear
}

func (p *EmailProcessor) extractEnmimeAttachments(env *enmime.Envelope, relativePath string) []Attachment {
	var attachments []Attachment
	outputDir := filepath.Join(p.config.EmailDir, relativePath)
	uuid := uuid.New().String()

	// Ensure output directory exists
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		log.Printf("Failed to create output directory: %v", err)
		return attachments
	}

	biggestHTMLSize := int64(0)
	biggestHTMLIdx := -1
	biggestPlainSize := int64(0)
	biggestPlainIdx := -1

	// Process HTML part if exists
	if env.HTML != "" {
		htmlData := []byte(env.HTML)
		fileName := "texthtml"
		newFileName := uuid + "-" + fileName
		filePath := filepath.Join(outputDir, newFileName)
		
		if err := ioutil.WriteFile(filePath, htmlData, 0644); err == nil {
			att := Attachment{
				Type:          "text/html",
				Name:          filepath.Join(relativePath, newFileName),
				Filename:      fileName,
				Size:          int64(len(htmlData)),
				IsMessageBody: 0,
			}
			attachments = append(attachments, att)
			biggestHTMLIdx = len(attachments) - 1
			biggestHTMLSize = int64(len(htmlData))
		}
	}

	// Process plain text part if exists
	if env.Text != "" {
		textData := []byte(env.Text)
		fileName := "textplain"
		newFileName := uuid + "-" + fileName
		filePath := filepath.Join(outputDir, newFileName)
		
		if err := ioutil.WriteFile(filePath, textData, 0644); err == nil {
			att := Attachment{
				Type:          "text/plain",
				Name:          filepath.Join(relativePath, newFileName),
				Filename:      fileName,
				Size:          int64(len(textData)),
				IsMessageBody: 0,
			}
			attachments = append(attachments, att)
			biggestPlainIdx = len(attachments) - 1
			biggestPlainSize = int64(len(textData))
		}
	}

	// Process file attachments
	for _, part := range env.Attachments {
		fileName := part.FileName
		if fileName == "" {
			fileName = "attachment"
		}
		
		newFileName := uuid + "-" + fileName
		filePath := filepath.Join(outputDir, newFileName)
		
		if err := ioutil.WriteFile(filePath, part.Content, 0644); err != nil {
			log.Printf("Failed to save attachment %s: %v", fileName, err)
			continue
		}

		att := Attachment{
			Type:          part.ContentType,
			Name:          filepath.Join(relativePath, newFileName),
			Filename:      fileName,
			Size:          int64(len(part.Content)),
			IsMessageBody: 0,
		}

		attachments = append(attachments, att)
		idx := len(attachments) - 1

		// Track largest HTML and plain text attachments
		if strings.HasPrefix(part.ContentType, "text/html") && int64(len(part.Content)) > biggestHTMLSize {
			biggestHTMLSize = int64(len(part.Content))
			biggestHTMLIdx = idx
		} else if strings.HasPrefix(part.ContentType, "text/plain") && int64(len(part.Content)) > biggestPlainSize {
			biggestPlainSize = int64(len(part.Content))
			biggestPlainIdx = idx
		}
	}

	// Process inline parts
	for _, part := range env.Inlines {
		fileName := part.FileName
		if fileName == "" {
			fileName = "inline"
		}
		
		newFileName := uuid + "-" + fileName
		filePath := filepath.Join(outputDir, newFileName)
		
		if err := ioutil.WriteFile(filePath, part.Content, 0644); err != nil {
			log.Printf("Failed to save inline part %s: %v", fileName, err)
			continue
		}

		att := Attachment{
			Type:          part.ContentType,
			Name:          filepath.Join(relativePath, newFileName),
			Filename:      fileName,
			Size:          int64(len(part.Content)),
			IsMessageBody: 0,
		}

		attachments = append(attachments, att)
		idx := len(attachments) - 1

		// Track largest HTML and plain text attachments
		if strings.HasPrefix(part.ContentType, "text/html") && int64(len(part.Content)) > biggestHTMLSize {
			biggestHTMLSize = int64(len(part.Content))
			biggestHTMLIdx = idx
		} else if strings.HasPrefix(part.ContentType, "text/plain") && int64(len(part.Content)) > biggestPlainSize {
			biggestPlainSize = int64(len(part.Content))
			biggestPlainIdx = idx
		}
	}

	// Mark the message body
	if biggestHTMLIdx >= 0 {
		attachments[biggestHTMLIdx].IsMessageBody = 1
	} else if biggestPlainIdx >= 0 {
		attachments[biggestPlainIdx].IsMessageBody = 1
	}

	return attachments
}

func (p *EmailProcessor) saveEmail(filename, subject string, unixTime int64, 
	recipientIDs map[string][]int, attachments []Attachment,
	foundEnquiries, foundInvoices, foundPOs, foundJobs []int) error {

	tx, err := p.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// Insert email
	result, err := tx.Exec(
		"INSERT INTO emails (user_id, udate, created, subject, filename) VALUES (?, ?, NOW(), ?, ?)",
		recipientIDs["from"][0], unixTime, subject, filename)
	
	if err != nil {
		return err
	}

	emailID, err := result.LastInsertId()
	if err != nil {
		return err
	}

	// Insert recipients
	for recipType, userIDs := range recipientIDs {
		for _, userID := range userIDs {
			if recipType == "from" {
				continue // From is already stored in emails.user_id
			}
			_, err = tx.Exec(
				"INSERT INTO email_recipients (email_id, user_id, type) VALUES (?, ?, ?)",
				emailID, userID, recipType)
			if err != nil {
				return err
			}
		}
	}

	// Insert attachments
	for _, att := range attachments {
		_, err = tx.Exec(
			"INSERT INTO email_attachments (email_id, name, type, size, filename, is_message_body, created) VALUES (?, ?, ?, ?, ?, ?, NOW())",
			emailID, att.Name, att.Type, att.Size, att.Filename, att.IsMessageBody)
		if err != nil {
			return err
		}
	}

	// Insert associations
	for _, jobID := range foundJobs {
		_, err = tx.Exec("INSERT INTO emails_jobs (email_id, job_id) VALUES (?, ?)", emailID, jobID)
		if err != nil {
			return err
		}
	}

	for _, poID := range foundPOs {
		_, err = tx.Exec("INSERT INTO emails_purchase_orders (email_id, purchase_order_id) VALUES (?, ?)", emailID, poID)
		if err != nil {
			return err
		}
	}

	for _, enqID := range foundEnquiries {
		_, err = tx.Exec("INSERT INTO emails_enquiries (email_id, enquiry_id) VALUES (?, ?)", emailID, enqID)
		if err != nil {
			return err
		}
	}

	for _, invID := range foundInvoices {
		_, err = tx.Exec("INSERT INTO emails_invoices (email_id, invoice_id) VALUES (?, ?)", emailID, invID)
		if err != nil {
			return err
		}
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	fmt.Println("Success. We made an email")
	return nil
}

func (p *EmailProcessor) moveEmail(filename string) error {
	oldPath := filepath.Join(p.config.VaultDir, filename)
	newPath := filepath.Join(p.config.ProcessedDir, filename+":S")

	if err := os.Rename(oldPath, newPath); err != nil {
		fmt.Printf("Unable to move %s to %s: %v\n", oldPath, newPath, err)
		return err
	}

	return nil
}