package main

import (
	"flag"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync"

	"github.com/schollz/progressbar/v3"
	"github.com/wholetrans/d10n/config"
	"github.com/wholetrans/d10n/translator"
)

func main() {
	// Parse command line arguments
	if len(os.Args) < 2 {
		printUsage()
		os.Exit(1)
	}

	sourcePath := os.Args[1]
	
	// Set up flags
	flagSet := flag.NewFlagSet("d10n", flag.ExitOnError)
	
	targetPathPtr := flagSet.String("target", "", "Target path for translated content")
	sourceLanguagePtr := flagSet.String("source-lang", "", "Source language code")
	targetLanguagePtr := flagSet.String("target-lang", "", "Target language code (required)")
	modelPtr := flagSet.String("model", "", "Model to use for translation")
	apiKeyPtr := flagSet.String("api-key", "", "API key for OpenAI compatible service")
	apiBasePtr := flagSet.String("api-base", "", "API base URL for OpenAI compatible service")
	systemPromptPtr := flagSet.String("system-prompt", "", "System prompt for the model")
	formatPtr := flagSet.String("format", "", "File format to process (e.g., md, txt)")
	
	// Concurrency flag
	concurrencyPtr := flagSet.Int("concurrency", 0, "Number of concurrent translation tasks (default is 3)")
	
	// Chunking flags
	chunkEnabledPtr := flagSet.Bool("chunk", false, "Enable chunked translation")
	chunkSizePtr := flagSet.Int("chunk-size", 0, "Size of each chunk in tokens (default is 10240)")
	chunkPromptPtr := flagSet.String("chunk-prompt", "", "Prompt to use for continuing translation (default is 'Please continue translation')")
	chunkContextPtr := flagSet.Int("chunk-context", 0, "Number of chunks to include as context (default is 2)")
	
	// Parse flags
	if err := flagSet.Parse(os.Args[2:]); err != nil {
		fmt.Println("Error parsing arguments:", err)
		os.Exit(1)
	}
	
	// Load configuration
	cfg, err := config.LoadConfig()
	if err != nil {
		fmt.Println("Error loading configuration:", err)
		os.Exit(1)
	}
	
	// Validate and set defaults
	if *targetLanguagePtr == "" {
		fmt.Println("Error: Target language is required")
		printUsage()
		os.Exit(1)
	}
	
	if *sourceLanguagePtr == "" {
		fmt.Println("Error: Source language is required")
		printUsage()
		os.Exit(1)
	}
	
	// Override config with command line arguments if provided
	if *modelPtr != "" {
		cfg.Model = *modelPtr
	}
	
	if *apiKeyPtr != "" {
		cfg.APIKey = *apiKeyPtr
	}
	
	if *apiBasePtr != "" {
		cfg.APIBase = *apiBasePtr
	}
	
	if *systemPromptPtr != "" {
		cfg.SystemPrompt = *systemPromptPtr
	}
	
	// Set concurrency options
	if *concurrencyPtr > 0 {
		cfg.Concurrency = *concurrencyPtr
	}
	
	// Set chunking options from command line arguments
	if *chunkEnabledPtr {
		cfg.Chunk.Enabled = true
	}
	
	if *chunkSizePtr > 0 {
		cfg.Chunk.Size = *chunkSizePtr
	}
	
	if *chunkPromptPtr != "" {
		cfg.Chunk.Prompt = *chunkPromptPtr
	}
	
	if *chunkContextPtr > 0 {
		cfg.Chunk.Context = *chunkContextPtr
	}
	
	// Set target path if not provided
	targetPath := *targetPathPtr
	if targetPath == "" {
		ext := filepath.Ext(sourcePath)
		base := strings.TrimSuffix(sourcePath, ext)
		targetPath = base + "_l10n" + ext
	}
	
	// Create translator
	trans := translator.NewTranslator(cfg)
	
	// Process the source path
	fileInfo, err := os.Stat(sourcePath)
	if err != nil {
		fmt.Printf("Error accessing source path %s: %v\n", sourcePath, err)
		os.Exit(1)
	}
	
	if fileInfo.IsDir() {
		err = processDirectory(sourcePath, targetPath, *sourceLanguagePtr, *targetLanguagePtr, *formatPtr, trans)
	} else {
		err = processFile(sourcePath, targetPath, *sourceLanguagePtr, *targetLanguagePtr, trans)
	}
	
	if err != nil {
		fmt.Println("Error during processing:", err)
		os.Exit(1)
	}
	
	fmt.Println("Translation completed successfully!")
}

func printUsage() {
	fmt.Println("Usage: d10n <source_path> [options]")
	fmt.Println("Options:")
	fmt.Println("  -target <path>          Target path for translated content")
	fmt.Println("  -source-lang <code>     Source language code (required)")
	fmt.Println("  -target-lang <code>     Target language code (required)")
	fmt.Println("  -model <model>          Model to use for translation")
	fmt.Println("  -api-key <key>          API key for OpenAI compatible service")
	fmt.Println("  -api-base <url>         API base URL for OpenAI compatible service")
	fmt.Println("  -system-prompt <prompt> System prompt for the model")
	fmt.Println("  -format <ext>           File format to process (e.g., md, txt)")
	fmt.Println("  -concurrency <num>      Number of concurrent translation tasks (default: 3)")
	fmt.Println("  -chunk                  Enable chunked translation")
	fmt.Println("  -chunk-size <tokens>    Size of each chunk in tokens (default: 10240)")
	fmt.Println("  -chunk-prompt <prompt>  Prompt for continuing translation (default: 'Please continue translation')")
	fmt.Println("  -chunk-context <num>    Number of chunks to include as context (default: 2)")
}

func processFile(sourcePath, targetPath, sourceLanguage, targetLanguage string, trans *translator.Translator) error {
	content, err := os.ReadFile(sourcePath)
	if err != nil {
		return fmt.Errorf("error reading file %s: %w", sourcePath, err)
	}
	
	translatedContent, err := trans.Translate(string(content), sourceLanguage, targetLanguage, sourcePath)
	if err != nil {
		return fmt.Errorf("error translating file %s: %w", sourcePath, err)
	}
	
	// Create target directory if it doesn't exist
	targetDir := filepath.Dir(targetPath)
	if err := os.MkdirAll(targetDir, 0755); err != nil {
		return fmt.Errorf("error creating target directory %s: %w", targetDir, err)
	}
	
	// Write translated content to target path
	if err := os.WriteFile(targetPath, []byte(translatedContent), 0644); err != nil {
		return fmt.Errorf("error writing to file %s: %w", targetPath, err)
	}
	
	// Remove the progress bar for this file
	trans.RemoveProgressBar(sourcePath)
	
	return nil
}

func processDirectory(sourcePath, targetPath, sourceLanguage, targetLanguage, format string, trans *translator.Translator) error {
	// Create target directory if it doesn't exist
	if err := os.MkdirAll(targetPath, 0755); err != nil {
		return fmt.Errorf("error creating target directory %s: %w", targetPath, err)
	}
	
	// First scan to find all matching files
	var matchedFiles []string
	err := filepath.Walk(sourcePath, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		
		// Skip directories
		if info.IsDir() {
			return nil
		}
		
		// Check if file matches the format
		if format != "" && !strings.HasSuffix(info.Name(), "."+format) {
			return nil
		}
		
		matchedFiles = append(matchedFiles, path)
		return nil
	})
	
	if err != nil {
		return err
	}
	
	// Display all matched files
	fmt.Printf("Found %d files to translate:\n", len(matchedFiles))
	for _, file := range matchedFiles {
		relPath, _ := filepath.Rel(sourcePath, file)
		fmt.Printf("- %s\n", relPath)
	}
	fmt.Println()
	
	// Create overall progress bar
	overallBar := progressbar.NewOptions(len(matchedFiles),
		progressbar.OptionSetDescription("[Overall Progress]"),
		progressbar.OptionShowCount(),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "#",
			SaucerHead:    ">",
			SaucerPadding: "-",
			BarStart:      "[",
			BarEnd:        "]",
		}),
	)
	
	// Set up a worker pool for concurrent processing
	var (
		wg          sync.WaitGroup
		mutex       sync.Mutex
		errors      []error
		concurrency = trans.GetConcurrency()
		semaphore   = make(chan struct{}, concurrency)
	)
	
	// Process files concurrently
	for i, path := range matchedFiles {
		// Compute relative path from source
		relPath, err := filepath.Rel(sourcePath, path)
		if err != nil {
			return fmt.Errorf("error computing relative path for %s: %w", path, err)
		}
		
		// Compute target file path
		targetFilePath := filepath.Join(targetPath, relPath)
		
		// Create target directory structure if needed
		targetFileDir := filepath.Dir(targetFilePath)
		if err := os.MkdirAll(targetFileDir, 0755); err != nil {
			return fmt.Errorf("error creating target directory %s: %w", targetFileDir, err)
		}
		
		// Process individual file concurrently
		wg.Add(1)
		go func(idx int, sourcePath, targetPath, relPath string) {
			defer wg.Done()
			
			// Acquire semaphore (limit concurrency)
			semaphore <- struct{}{}
			defer func() { <-semaphore }()
			
			// Process the file
			err := processFile(sourcePath, targetPath, sourceLanguage, targetLanguage, trans)
			
			// Update overall progress bar
			mutex.Lock()
			overallBar.Add(1)
			if err != nil {
				errors = append(errors, err)
				fmt.Printf("Error translating %s: %v\n", relPath, err)
			}
			mutex.Unlock()
		}(i+1, path, targetFilePath, relPath)
	}
	
	// Wait for all translations to complete
	wg.Wait()
	
	// Complete the overall progress bar
	overallBar.Finish()
	fmt.Println() // Add some spacing after progress bars
	
	// Check if any errors occurred
	if len(errors) > 0 {
		return fmt.Errorf("encountered %d errors during translation", len(errors))
	}
	
	return nil
}