package main

import (
	"context"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"

	"golang.org/x/oauth2"
	"golang.org/x/oauth2/google"
	"google.golang.org/api/gmail/v1"
	"google.golang.org/api/option"
)

type EmailInfo struct {
	Sender  string `json:"sender"`
	Subject string `json:"subject"`
	Date    string `json:"date"`
	ID      string `json:"id"`
}

func main() {
	if len(os.Args) < 4 {
		fmt.Println("Usage: go run main.go <client_id> <client_secret> <output_format>")
		fmt.Println("  client_id: Your Google OAuth2 client ID")
		fmt.Println("  client_secret: Your Google OAuth2 client secret")
		fmt.Println("  output_format: csv or json")
		os.Exit(1)
	}

	clientID := os.Args[1]
	clientSecret := os.Args[2]
	outputFormat := strings.ToLower(os.Args[3])
	
	if outputFormat != "csv" && outputFormat != "json" {
		log.Fatal("Output format must be 'csv' or 'json'")
	}

	ctx := context.Background()

	config := &oauth2.Config{
		ClientID:     clientID,
		ClientSecret: clientSecret,
		RedirectURL:  "urn:ietf:wg:oauth:2.0:oob",
		Scopes:       []string{gmail.GmailReadonlyScope},
		Endpoint:     google.Endpoint,
	}

	client := getClient(config)

	srv, err := gmail.NewService(ctx, option.WithHTTPClient(client))
	if err != nil {
		log.Fatalf("Unable to retrieve Gmail client: %v", err)
	}

	fmt.Printf("Connected to Gmail API successfully\n")
	fmt.Printf("Fetching emails from inbox...\n")
	
	emails := fetchEmails(srv)

	fmt.Printf("Generating %s output...\n", outputFormat)
	
	switch outputFormat {
	case "csv":
		outputCSV(emails)
	case "json":
		outputJSON(emails)
	}

	fmt.Printf("Complete! Fetched %d emails\n", len(emails))
}

func getClient(config *oauth2.Config) *http.Client {
	authURL := config.AuthCodeURL("state-token", oauth2.AccessTypeOffline)
	fmt.Printf("Visit this URL in your browser: %v\n\n", authURL)
	fmt.Print("Paste the authorization code here: ")

	var authCode string
	if _, err := fmt.Scanln(&authCode); err != nil {
		log.Fatalf("Unable to read authorization code: %v", err)
	}

	tok, err := config.Exchange(context.TODO(), authCode)
	if err != nil {
		log.Fatalf("Unable to retrieve token: %v", err)
	}
	
	return config.Client(context.Background(), tok)
}

func fetchEmails(srv *gmail.Service) []EmailInfo {
	var emails []EmailInfo
	pageCount := 0
	
	req := srv.Users.Messages.List("me").Q("in:inbox")
	
	for {
		pageCount++
		fmt.Printf("Fetching page %d...\n", pageCount)
		
		r, err := req.Do()
		if err != nil {
			log.Fatalf("Unable to retrieve messages: %v", err)
		}

		fmt.Printf("Processing %d messages from page %d...\n", len(r.Messages), pageCount)
		
		for i, m := range r.Messages {
			if (i+1)%50 == 0 {
				fmt.Printf("  Processed %d/%d messages on this page\n", i+1, len(r.Messages))
			}
			
			msg, err := srv.Users.Messages.Get("me", m.Id).Do()
			if err != nil {
				log.Printf("Unable to retrieve message %s: %v", m.Id, err)
				continue
			}

			email := parseMessage(msg)
			emails = append(emails, email)
		}

		fmt.Printf("Completed page %d. Total emails collected: %d\n", pageCount, len(emails))

		if r.NextPageToken == "" {
			fmt.Printf("Reached end of inbox. Processing complete.\n")
			break
		}
		req.PageToken(r.NextPageToken)
	}

	return emails
}

func parseMessage(msg *gmail.Message) EmailInfo {
	email := EmailInfo{
		ID: msg.Id,
	}

	for _, header := range msg.Payload.Headers {
		switch header.Name {
		case "From":
			email.Sender = header.Value
		case "Subject":
			email.Subject = header.Value
		case "Date":
			email.Date = header.Value
		}
	}

	return email
}

func outputCSV(emails []EmailInfo) {
	writer := csv.NewWriter(os.Stdout)
	defer writer.Flush()

	writer.Write([]string{"Sender", "Subject", "Date", "ID"})

	for _, email := range emails {
		writer.Write([]string{email.Sender, email.Subject, email.Date, email.ID})
	}
}

func outputJSON(emails []EmailInfo) {
	encoder := json.NewEncoder(os.Stdout)
	encoder.SetIndent("", "  ")
	encoder.Encode(emails)
}