internal/generator/generator.go

package generator

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"io/fs"
	"log"
	"os"
	"path"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"static-repo/internal/models"
	"static-repo/internal/renderer"
	"static-repo/internal/tree"
	"static-repo/internal/utils"
)

type Config struct {
	ContentDir string
	OutputDir  string
	LayoutFile string
	Subdir     string
}

type Generator struct {
	config Config
	rend   *renderer.Renderer
	tmpl   *template.Template
	fs     fs.FS
}

func NewGenerator(config Config, templatesFS fs.FS) (*Generator, error) {
	rend, err := renderer.NewRenderer(templatesFS)
	if err != nil {
		return nil, fmt.Errorf("failed to create renderer: %v", err)
	}

	tmpl, err := template.New(filepath.Base(config.LayoutFile)).Funcs(template.FuncMap{
		"dict": func(values ...interface{}) (map[string]interface{}, error) {
			if len(values)%2 != 0 {
				return nil, fmt.Errorf("invalid dict call")
			}
			dict := make(map[string]interface{}, len(values)/2)
			for i := 0; i < len(values); i += 2 {
				key, ok := values[i].(string)
				if !ok {
					return nil, fmt.Errorf("dict keys must be strings")
				}
				dict[key] = values[i+1]
			}
			return dict, nil
		},
		"isAncestor": func(current, nodePath string) bool {
			if nodePath == "" || nodePath == "." || nodePath == "Root" {
				return true
			}
			current = filepath.ToSlash(current)
			nodePath = filepath.ToSlash(nodePath)
			return current == nodePath || strings.HasPrefix(current, nodePath+"/")
		},
	}).ParseFS(templatesFS, config.LayoutFile)
	if err != nil {
		return nil, fmt.Errorf("failed to parse template: %v", err)
	}

	// Parse partials
	tmpl, err = tmpl.ParseFS(templatesFS, "templates/partials/*.html")
	if err != nil {
		return nil, fmt.Errorf("failed to parse partials: %v", err)
	}

	// Parse icons
	tmpl, err = tmpl.ParseFS(templatesFS, "templates/partials/icons/*.svg")
	if err != nil {
		return nil, fmt.Errorf("failed to parse icons: %v", err)
	}

	// Parse projects index template
	tmpl, err = tmpl.ParseFS(templatesFS, "templates/projects_index.html")
	if err != nil {
		return nil, fmt.Errorf("failed to parse projects index template: %v", err)
	}

	return &Generator{
		config: config,
		rend:   rend,
		tmpl:   tmpl,
		fs:     templatesFS,
	}, nil
}

var generationCount int

func (g *Generator) Generate() error {
	start := time.Now()

	// Ensure output directory exists and is clean
	if err := os.RemoveAll(g.config.OutputDir); err != nil {
		return fmt.Errorf("failed to clean output directory: %v", err)
	}
	if err := os.MkdirAll(g.config.OutputDir, 0755); err != nil {
		return fmt.Errorf("failed to create output directory: %v", err)
	}

	projects, err := g.collectProjects()
	defer func() {
		for _, p := range projects {
			if p.TempDir != "" {
				fmt.Printf("Cleaning up temp dir for project %s: %s\n", p.Title, p.TempDir)
				os.RemoveAll(p.TempDir)
			}
		}
	}()
	if err != nil {
		return err
	}

	for _, project := range projects {
		if err := g.processProject(project, projects); err != nil {
			log.Printf("Warning: failed to process project %s: %v", project.Name, err)
		}
	}

	if err := g.renderOverview(projects); err != nil {
		return err
	}

	if err := g.copyAssets(); err != nil {
		return err
	}

	fmt.Printf("Generated %d pages in %s\n", generationCount, time.Since(start))
	return nil
}

func (g *Generator) getBaseURL() string {
	prefix := "/"
	if g.config.Subdir != "" {
		prefix = path.Join(prefix, g.config.Subdir)
	}
	if !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	return prefix
}

func (g *Generator) getURLPrefix(projectName string) string {
	baseURL := g.getBaseURL()
	if projectName == "" {
		return baseURL
	}
	return path.Join(baseURL, projectName)
}

func (g *Generator) collectProjects() ([]models.Project, error) {
	var projects []models.Project

	// Process external.json if it exists
	externalPath := filepath.Join(g.config.ContentDir, "external.json")
	if _, err := os.Stat(externalPath); err == nil {
		data, err := os.ReadFile(externalPath)
		if err == nil {
			var extProjects []struct {
				URL         string `json:"url"`
				Title       string `json:"title"`
				Description string `json:"description"`
				Date        string `json:"date"`
			}
			if err := json.Unmarshal(data, &extProjects); err == nil {
				for _, ep := range extProjects {
					fmt.Printf("Downloading external project: %s\n", ep.Title)
					proj, err := g.downloadExternalProject(ep.URL, ep.Title, ep.Description, ep.Date)
					if err != nil {
						log.Printf("Warning: %v", err)
						continue
					}
					projects = append(projects, proj)
				}
			}
		}
	}

	rootEntries, err := os.ReadDir(g.config.ContentDir)
	if err != nil {
		return projects, fmt.Errorf("failed to read content directory: %v", err)
	}

	for _, entry := range rootEntries {
		if !entry.IsDir() {
			continue
		}
		projectName := entry.Name()
		projectDir := filepath.Join(g.config.ContentDir, projectName)
		urlPrefix := g.getURLPrefix(projectName)
		project := models.Project{
			Name:       projectName,
			Title:      projectName,
			ContentDir: projectDir,
			Link:       path.Join(urlPrefix, "index.html"),
			ZipLink:    path.Join(urlPrefix, projectName+".zip"),
		}

		readmePath := g.findReadme(projectDir)
		if readmePath != "" {
			meta, err := g.rend.ExtractFrontmatter(readmePath)
			if err == nil {
				if title, ok := meta["title"].(string); ok {
					project.Title = title
				}
				if desc, ok := meta["description"].(string); ok {
					project.Description = desc
				}
				project.Date = normalizeProjectDate(meta["date"])
			}
		}
		projects = append(projects, project)
	}

	return projects, nil
}

func (g *Generator) downloadExternalProject(url, title, description, date string) (models.Project, error) {
	zipFile, err := os.CreateTemp("", "ext-*.zip")
	if err != nil {
		return models.Project{}, fmt.Errorf("failed to create temp file: %v", err)
	}
	zipFilePath := zipFile.Name()
	zipFile.Close()
	defer os.Remove(zipFilePath)

	if err := utils.DownloadFile(url, zipFilePath); err != nil {
		return models.Project{}, fmt.Errorf("failed to download %s: %v", url, err)
	}

	tempDir, err := os.MkdirTemp("", "ext-project-*")
	if err != nil {
		return models.Project{}, fmt.Errorf("failed to create temp dir: %v", err)
	}

	if err := utils.Unzip(zipFilePath, tempDir); err != nil {
		os.RemoveAll(tempDir)
		return models.Project{}, fmt.Errorf("failed to unzip: %v", err)
	}

	entries, _ := os.ReadDir(tempDir)
	actualContentDir := tempDir
	if len(entries) == 1 && entries[0].IsDir() {
		actualContentDir = filepath.Join(tempDir, entries[0].Name())
	}

	projectName := strings.ToLower(strings.ReplaceAll(title, " ", "-"))
	urlPrefix := g.getURLPrefix(projectName)
	return models.Project{
		Name:        projectName,
		Title:       title,
		Description: description,
		Date:        date,
		ContentDir:  actualContentDir,
		Link:        path.Join(urlPrefix, "index.html"),
		TempDir:     tempDir,
		ZipLink:     path.Join(urlPrefix, projectName+".zip"),
	}, nil
}

func normalizeProjectDate(value any) string {
	switch v := value.(type) {
	case string:
		return strings.TrimSpace(v)
	case time.Time:
		return v.Format("2006-01-02")
	default:
		return ""
	}
}

func (g *Generator) processProject(project models.Project, allProjects []models.Project) error {
	projectOutputDir := filepath.Join(g.config.OutputDir, project.Name)
	zipName := project.Name + ".zip"

	if err := os.MkdirAll(projectOutputDir, 0755); err != nil {
		return err
	}

	fileTree, err := tree.BuildFileTree(project.ContentDir, "", g.getURLPrefix(project.Name))
	if err != nil {
		return err
	}

	// Collect files to render
	var files []string
	err = filepath.WalkDir(project.ContentDir, func(filePath string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.IsDir() {
			files = append(files, filePath)
		}
		return nil
	})
	if err != nil {
		return err
	}

	numFiles := len(files)
	if numFiles == 0 {
		return nil
	}

	// Use worker pool for parallel rendering
	numWorkers := runtime.NumCPU()
	if numWorkers > numFiles {
		numWorkers = numFiles
	}

	jobs := make(chan string, numFiles)
	results := make(chan error, numFiles)
	var wg sync.WaitGroup

	for w := 0; w < numWorkers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for filePath := range jobs {
				relPath, _ := filepath.Rel(project.ContentDir, filePath)
				outRel := relPath + ".html"
				if strings.ToLower(filepath.Base(relPath)) == "readme.md" {
					outRel = filepath.Join(filepath.Dir(relPath), "index.html")
				}
				outputPath := filepath.Join(projectOutputDir, outRel)
				os.MkdirAll(filepath.Dir(outputPath), 0755)

				urlPrefix := g.getURLPrefix(project.Name)
				content, err := g.rend.RenderFile(filePath, project.ContentDir, projectOutputDir, urlPrefix)
				if err != nil {
					log.Printf("Warning: failed to render %s: %v", filePath, err)
					results <- nil // Continue on render error
					continue
				}

				// Copy the original file as well, so users can download it
				if strings.HasSuffix(strings.ToLower(filePath), ".md") {
					origPath := filepath.Join(projectOutputDir, relPath)
					os.MkdirAll(filepath.Dir(origPath), 0755)
					utils.CopyFile(filePath, origPath)
				}

				generationCount++

				data := models.TemplateData{
					Title:          filepath.ToSlash(relPath),
					Content:        template.HTML(content),
					FileTree:       fileTree,
					CurrentProject: &project,
					Projects:       allProjects,
					CurrentPath:    filepath.ToSlash(relPath),
					BaseURL:        g.getBaseURL(),
				}

				f, err := os.Create(outputPath)
				if err != nil {
					results <- err
					continue
				}

				if err := g.tmpl.Execute(f, data); err != nil {
					f.Close()
					results <- err
					continue
				}
				f.Close()
				results <- nil
			}
		}()
	}

	for _, path := range files {
		jobs <- path
	}
	close(jobs)

	go func() {
		wg.Wait()
		close(results)
	}()

	for err := range results {
		if err != nil {
			return err
		}
	}

	// Render project index
	readmePath := ""
	currentPath := ""
	for _, child := range fileTree.Children {
		if strings.ToLower(child.Name) == "readme.md" {
			readmePath = filepath.Join(project.ContentDir, child.Path)
			currentPath = child.Path
			break
		}
	}

	var content template.HTML
	if readmePath != "" {
		c, err := g.rend.RenderMarkdown(readmePath)
		if err != nil {
			return err
		}
		content = template.HTML(c)
	}

	indexPath := filepath.Join(projectOutputDir, "index.html")
	indexFile, err := os.Create(indexPath)
	if err != nil {
		return err
	}
	defer indexFile.Close()

	data := models.TemplateData{
		Title:          project.Title,
		Content:        content,
		FileTree:       fileTree,
		CurrentProject: &project,
		Projects:       allProjects,
		CurrentPath:    currentPath,
		BaseURL:        g.getBaseURL(),
		IsIndex:        true,
	}
	if err := g.tmpl.Execute(indexFile, data); err != nil {
		return err
	}

	return utils.ZipDirectory(project.ContentDir, filepath.Join(projectOutputDir, zipName))
}

func (g *Generator) renderOverview(projects []models.Project) error {
	overviewPath := filepath.Join(g.config.OutputDir, "index.html")
	f, err := os.Create(overviewPath)
	if err != nil {
		return err
	}
	defer f.Close()

	overviewTitle := "Project Overview"
	var rootReadmeContent template.HTML
	rootReadmePath := g.findReadme(g.config.ContentDir)
	if rootReadmePath != "" {
		meta, err := g.rend.ExtractFrontmatter(rootReadmePath)
		if err == nil {
			if title, ok := meta["title"].(string); ok {
				overviewTitle = title
			}
		}
		content, err := g.rend.RenderMarkdown(rootReadmePath)
		if err == nil {
			rootReadmeContent = template.HTML(content)
		}
	}

	var contentBuf bytes.Buffer
	if err := g.tmpl.ExecuteTemplate(&contentBuf, "projects_index.html", models.TemplateData{
		Projects: projects,
		Extra:    map[string]any{"OverviewContent": rootReadmeContent},
	}); err != nil {
		return err
	}

	data := models.TemplateData{
		Title:    overviewTitle,
		Content:  template.HTML(contentBuf.String()),
		Projects: projects,
		BaseURL:  g.getBaseURL(),
		IsIndex:  true,
	}

	return g.tmpl.Execute(f, data)
}

func (g *Generator) copyAssets() error {
	assetsDir := filepath.Join(g.config.OutputDir, "assets")
	os.MkdirAll(filepath.Join(assetsDir, "icons"), 0755)

	style, err := fs.ReadFile(g.fs, "templates/style.min.css")
	if err == nil {
		os.WriteFile(filepath.Join(assetsDir, "style.min.css"), style, 0644)
	}

	script, err := fs.ReadFile(g.fs, "templates/script.js")
	if err == nil {
		os.WriteFile(filepath.Join(assetsDir, "script.js"), script, 0644)
	}

	iconEntries, _ := fs.ReadDir(g.fs, "templates/partials/icons")
	for _, entry := range iconEntries {
		data, err := fs.ReadFile(g.fs, filepath.Join("templates/partials/icons", entry.Name()))
		if err == nil {
			dst := filepath.Join(assetsDir, "icons", entry.Name())
			os.WriteFile(dst, data, 0644)
		}
	}

	icon, err := os.Open(filepath.Join(g.config.ContentDir, "favicon.ico"))
	if err == nil {
		f, _ := os.Create(filepath.Join(g.config.OutputDir, "favicon.ico"))
		io.Copy(f, icon)
		f.Close()
	}

	return nil
}

func (g *Generator) findReadme(dir string) string {
	entries, _ := os.ReadDir(dir)
	for _, e := range entries {
		if strings.ToLower(e.Name()) == "readme.md" {
			return filepath.Join(dir, e.Name())
		}
	}
	return ""
}