package main

import (
	"archive/zip"
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"maps"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
	"time"

	"go.mau.fi/whatsmeow/store/sqlstore"
	waLog "go.mau.fi/whatsmeow/util/log"
)

const (
	// whatsappImportMark marks history in the first line of a batch.
	whatsappImportMark = " (imported)"

	// whatsappImportBytes limits one export's uncompressed chat text.
	whatsappImportBytes = 64 << 20
)

var (
	whatsappImportStart   = regexp.MustCompile(`^\[(\d{2}/\d{2}/\d{4}, \d{2}:\d{2}:\d{2})\] ([^\r\n]+?): ?(.*)$`)
	whatsappImportHeader  = regexp.MustCompile(`^\[WhatsApp\] (.*), (?:group|private chat) ([^ ,]+), [A-Za-z]+ (\d{4}-\d{2}-\d{2})$`)
	whatsappImportUnknown = regexp.MustCompile(`^\[WhatsApp\] (.*), [A-Za-z]+ (\d{4}-\d{2}-\d{2})$`)
)

// whatsappImport runs the offline import command. It holds the same token lock
// as the daemon and probe, even for a plan. The plan reads no credentials for a
// model and makes no network call. Only a completed rebuild replaces the agent
// directory, and unfinished subagents prevent either mode.
func whatsappImport(ctx context.Context, args []string, out io.Writer) (err error) {
	planOnly := len(args) > 0 && args[0] == "-n"
	if planOnly {
		args = args[1:]
	}

	if len(args) == 0 {
		return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
	}

	for _, arg := range args {
		if strings.HasPrefix(arg, "-") {
			return errors.New("usage: onechat import [-n] ZIP [ZIP...]")
		}
	}

	dir, err := dataDir()
	if err != nil {
		return fmt.Errorf("find import data directory: %w", err)
	}

	tokens, err := lockTokens(dir)
	if err != nil {
		return fmt.Errorf("lock import tokens: %w", err)
	}

	defer func() {
		closeErr := tokens.Close()
		if closeErr != nil {
			err = errors.Join(err, fmt.Errorf("close import token lock: %w", closeErr))
		}
	}()

	root := filepath.Join(dir, "agent")
	err = rebuildIdle(root)
	if err != nil {
		return fmt.Errorf("check import work: %w", err)
	}

	old, err := rebuildReadLog(ctx, root)
	if err != nil {
		return fmt.Errorf("read import log: %w", err)
	}

	names, owner, err := whatsappImportNames(ctx, dir)
	if err != nil {
		return fmt.Errorf("read import names: %w", err)
	}

	plan, err := whatsappImportPlan(ctx, args, old, names, owner, time.Local)
	if err != nil {
		return fmt.Errorf("plan import: %w", err)
	}

	err = plan.Write(out)
	if err != nil {
		return fmt.Errorf("write import plan: %w", err)
	}

	if planOnly || len(plan.added) == 0 {
		return nil
	}

	// An existing agent is required so that a failed import never creates
	// or changes the original agent's configuration.
	_, err = os.Stat(filepath.Join(root, "config.json"))
	if err != nil {
		return fmt.Errorf("read agent config: %w", err)
	}

	agent, err := openAgent(root)
	if err != nil {
		return fmt.Errorf("open import agent: %w", err)
	}

	config, err := loadConfig(dir)
	if err != nil {
		return fmt.Errorf("load import configuration: %w", err)
	}

	address, auth, err := modelEndpoint(config, dir)
	if err != nil {
		return fmt.Errorf("select import endpoint: %w", err)
	}

	usage, err := os.CreateTemp(dir, ".import-usage-*.jsonl")
	if err != nil {
		return fmt.Errorf("create import usage: %w", err)
	}

	usagePath := usage.Name()
	defer func() {
		closeErr := os.Remove(usagePath)
		if closeErr != nil {
			err = errors.Join(err, fmt.Errorf("remove import usage: %w", closeErr))
		}
	}()

	err = usage.Close()
	if err != nil {
		return fmt.Errorf("close import usage: %w", err)
	}

	client := newModelClient(address, auth, agent.CacheKey(), usagePath)
	backup, err := rebuild(ctx, agent, plan.added, client, newCompactor(client, agent.Prompt, config.Compactor), out)
	if err != nil {
		return fmt.Errorf("rebuild imported memory: %w", err)
	}

	_, err = fmt.Fprintf(out, "import complete: backup %s\n", backup)
	if err != nil {
		return fmt.Errorf("write import result: %w", err)
	}

	return nil
}

// WhatsAppImportPlan holds the new chat-day batches and the size of the merged
// log. Its estimate uses the measured cost of 1.3 calls per message, not a
// promise about a particular export or compaction retry.
type WhatsAppImportPlan struct {
	added      []DatedMessage
	chats      []string
	cutoff     time.Time
	first      time.Time
	last       time.Time
	bytes      int
	messages   int
	skipped    int
	ownerKnown bool
}

// Write prints the plan without chat text or the owner's profile name.
func (p WhatsAppImportPlan) Write(out io.Writer) error {
	cutoff := "none (no live WhatsApp batch)"
	if !p.cutoff.IsZero() {
		cutoff = p.cutoff.Format(time.RFC3339)
	}

	dateRange := "none"
	if !p.first.IsZero() {
		dateRange = p.first.Format(time.RFC3339) + " to " + p.last.Format(time.RFC3339)
	}

	owner := "known (mapped to You)"
	if !p.ownerKnown {
		owner = "unknown (sender names kept as written)"
	}

	_, err := fmt.Fprintf(out, "chats: %d\n", len(p.chats))
	if err != nil {
		return fmt.Errorf("write import plan: %w", err)
	}

	for _, chat := range p.chats {
		_, err = fmt.Fprintf(out, "  %s\n", chat)
		if err != nil {
			return fmt.Errorf("write import plan: %w", err)
		}
	}

	estimate := 1.3 * float64(p.messages)
	if len(p.added) == 0 {
		estimate = 0
	}

	lines := []string{
		fmt.Sprintf("batches: %d (%d chat-days skipped)", len(p.added), p.skipped),
		"dates: " + dateRange,
		"cutoff: " + cutoff,
		"owner: " + owner,
		fmt.Sprintf("bytes: %d (new batch text)", p.bytes),
		fmt.Sprintf("messages: %d (merged log)", p.messages),
		fmt.Sprintf("estimated calls: %.0f (about 1.3 per message, limit %d)", estimate, rebuildCallLimit),
	}

	_, err = io.WriteString(out, strings.Join(lines, "\n")+"\n")
	if err != nil {
		return fmt.Errorf("write import plan: %w", err)
	}

	return nil
}

// whatsappImportNames reads only local session data. Group names that the
// library keeps only in memory can also come from logged or pending batches.
// Ambiguous names have no ID. A missing session keeps sender names unchanged.
func whatsappImportNames(ctx context.Context, dir string) (names map[string]WhatsAppChat, owner string, err error) {
	names = make(map[string]WhatsAppChat)
	path := filepath.Join(dir, whatsappDirName, whatsappSessionFileName)
	_, err = os.Stat(path)
	if errors.Is(err, os.ErrNotExist) {
		return names, "", nil
	}

	if err != nil {
		return nil, "", fmt.Errorf("read import session: %w", err)
	}

	address, err := whatsappSessionAddress(path)
	if err != nil {
		return nil, "", fmt.Errorf("find import session address: %w", err)
	}

	container, err := sqlstore.New(ctx, "sqlite", address, waLog.Noop)
	if err != nil {
		return nil, "", fmt.Errorf("open import session: %w", err)
	}

	defer func() {
		closeErr := container.Close()
		if closeErr != nil {
			err = errors.Join(err, fmt.Errorf("close import session: %w", closeErr))
		}
	}()

	device, err := container.GetFirstDevice(ctx)
	if err != nil {
		return nil, "", fmt.Errorf("read import device: %w", err)
	}

	if device.ID == nil {
		return names, "", nil
	}

	owner = device.PushName
	own := whatsappOwnIDs(device.ID.User, device.LID.User)
	contacts, err := device.Contacts.GetAllContacts(ctx)
	if err != nil {
		return nil, "", fmt.Errorf("read import contacts: %w", err)
	}

	for jid, contact := range contacts {
		if slices.Contains(own, jid.User) {
			continue
		}

		for _, name := range []string{contact.FirstName, contact.FullName, contact.PushName, contact.BusinessName} {
			if name != "" {
				whatsappImportName(names, WhatsAppChat{
					ID:    jid.ToNonAD().String(),
					Name:  name,
					Group: jid.Server == "g.us",
				})
			}
		}
	}

	data, err := os.ReadFile(filepath.Join(dir, whatsappDirName, whatsappStateFileName))
	if errors.Is(err, os.ErrNotExist) {
		return names, owner, nil
	}

	if err != nil {
		return nil, "", fmt.Errorf("read pending import names: %w", err)
	}

	var state WhatsAppStateFile

	err = json.Unmarshal(data, &state)
	if err != nil {
		return nil, "", fmt.Errorf("decode pending import names: %w", err)
	}

	for _, batch := range state.Batches {
		whatsappImportName(names, batch.Chat)
	}

	return names, owner, nil
}

// whatsappImportName adds an exact match unless distinct IDs share its name.
func whatsappImportName(names map[string]WhatsAppChat, chat WhatsAppChat) {
	old, found := names[chat.Name]
	if found && old.ID != chat.ID {
		chat.ID = ""
	}

	names[chat.Name] = chat
}

// whatsappImportPlan filters exports at the first message of the earliest live
// batch and skips chat-days already imported. Each remaining chat-day becomes a
// work message dated at its last exported message. Overlapping chat-days in the
// inputs are refused, not silently combined or duplicated.
func whatsappImportPlan(ctx context.Context, paths []string, old []DatedMessage, names map[string]WhatsAppChat, owner string, location *time.Location) (WhatsAppImportPlan, error) {
	plan := WhatsAppImportPlan{ownerKnown: owner != ""}
	names = maps.Clone(names)
	if names == nil {
		names = make(map[string]WhatsAppChat)
	}

	imported := make(map[string]bool)
	for _, entry := range old {
		if entry.Kind != "work" || !strings.HasPrefix(entry.Text, "[WhatsApp] ") {
			continue
		}

		first, _, _ := strings.Cut(entry.Text, "\n")
		marked := strings.HasSuffix(first, whatsappImportMark)
		first = strings.TrimSuffix(first, whatsappImportMark)
		match := whatsappImportHeader.FindStringSubmatch(first)
		name, day := "", ""
		if len(match) > 0 {
			name, day = match[1], match[3]
			whatsappImportName(names, WhatsAppChat{
				ID:    match[2],
				Name:  name,
				Group: strings.Contains(first, ", group "+match[2]+","),
			})
		} else if marked {
			match = whatsappImportUnknown.FindStringSubmatch(first)
			if len(match) > 0 {
				name, day = match[1], match[2]
			}
		}

		if name == "" {
			return plan, errors.New("read WhatsApp history: invalid batch header")
		}

		if marked {
			imported[name+"\n"+day] = true
			continue
		}

		_, body, found := strings.Cut(entry.Text, "\n")
		if !found || len(body) < len("15:04 ") {
			return plan, errors.New("read WhatsApp history: missing first message time")
		}

		date, err := time.ParseInLocation("2006-01-02 15:04", day+" "+body[:5], location)
		if err != nil {
			return plan, fmt.Errorf("parse first live message time: %w", err)
		}

		if plan.cutoff.IsZero() || date.Before(plan.cutoff) {
			plan.cutoff = date
		}
	}

	seen := make(map[string]bool)
	for _, path := range paths {
		if ctx.Err() != nil {
			return plan, ctx.Err()
		}

		chat, messages, err := whatsappImportRead(path, owner, location)
		if err != nil {
			return plan, fmt.Errorf("read import export: %w", err)
		}

		known, found := names[chat.Name]
		if found {
			chat = known
		}

		if !slices.Contains(plan.chats, chat.Name) {
			plan.chats = append(plan.chats, chat.Name)
		}

		days := make(map[string][]WhatsAppMessage)
		for _, message := range messages {
			if !plan.cutoff.IsZero() && !message.Time.Before(plan.cutoff) {
				continue
			}

			day := message.Time.In(location).Format(time.DateOnly)
			days[day] = append(days[day], message)
		}

		for _, day := range slices.Sorted(maps.Keys(days)) {
			key := chat.Name + "\n" + day
			if imported[key] {
				plan.skipped++
				continue
			}

			if seen[key] {
				return plan, errors.New("import inputs overlap: give only one export for each chat and day")
			}

			seen[key] = true
			batch := WhatsAppBatch{
				Chat:     chat,
				Messages: days[day],
			}

			text := batch.Text(location)
			first, body, _ := strings.Cut(text, "\n")
			if chat.ID == "" {
				first = "[WhatsApp] " + chat.Name + ", " + batch.Messages[0].Time.In(location).Format(whatsappBatchDayLayout)
			}

			text = first + whatsappImportMark + "\n" + body
			plan.added = append(plan.added, DatedMessage{
				Kind: "work",
				Text: text,
				Date: batch.Messages[len(batch.Messages)-1].Time.Format(time.RFC3339Nano),
			})

			plan.bytes += len(text)
			firstDate, lastDate := batch.Messages[0].Time, batch.Messages[len(batch.Messages)-1].Time
			if plan.first.IsZero() || firstDate.Before(plan.first) {
				plan.first = firstDate
			}

			if plan.last.IsZero() || lastDate.After(plan.last) {
				plan.last = lastDate
			}
		}
	}

	var err error

	plan.added, err = rebuildMerge(nil, plan.added)
	if err != nil {
		return plan, fmt.Errorf("order import batches: %w", err)
	}

	plan.messages = len(old) + len(plan.added)
	slices.Sort(plan.chats)

	return plan, nil
}

// whatsappImportRead reads the single _chat.txt in an export without extracting
// archive paths. Non-start lines continue the preceding message, and all
// invisible left-to-right marks are removed. Dates use location. Media and
// system messages remain text, including the written sender.
func whatsappImportRead(path, owner string, location *time.Location) (chat WhatsAppChat, messages []WhatsAppMessage, err error) {
	chat.Name = strings.TrimPrefix(strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)), "WhatsApp Chat - ")
	if chat.Name == "" || strings.ContainsAny(chat.Name, "\r\n") {
		return chat, nil, errors.New("read export: chat name must be one nonempty line")
	}

	archive, err := zip.OpenReader(path)
	if err != nil {
		return chat, nil, fmt.Errorf("open export: %w", err)
	}

	defer func() {
		closeErr := archive.Close()
		if closeErr != nil {
			err = errors.Join(err, fmt.Errorf("close export: %w", closeErr))
		}
	}()

	var file *zip.File

	for _, entry := range archive.File {
		if filepath.Base(entry.Name) == "_chat.txt" && !entry.FileInfo().IsDir() {
			if file != nil {
				return chat, nil, errors.New("read export: more than one _chat.txt")
			}

			file = entry
		}
	}

	if file == nil {
		return chat, nil, errors.New("read export: _chat.txt is missing")
	}

	if file.UncompressedSize64 > whatsappImportBytes {
		return chat, nil, errors.New("read export: chat text exceeds 64 MiB")
	}

	reader, err := file.Open()
	if err != nil {
		return chat, nil, fmt.Errorf("open chat text: %w", err)
	}

	defer func() {
		closeErr := reader.Close()
		if closeErr != nil {
			err = errors.Join(err, fmt.Errorf("close export chat text: %w", closeErr))
		}
	}()

	scanner := bufio.NewScanner(io.LimitReader(reader, whatsappImportBytes+1))
	scanner.Buffer(make([]byte, 64*1024), buildLineBytes)

	for scanner.Scan() {
		line := strings.ReplaceAll(scanner.Text(), "\u200e", "")
		match := whatsappImportStart.FindStringSubmatch(line)
		if match == nil {
			if len(messages) == 0 {
				if line == "" {
					continue
				}

				return chat, nil, errors.New("read export: text before the first dated message")
			}

			messages[len(messages)-1].Text += "\n" + line

			continue
		}

		var date time.Time

		date, err = time.ParseInLocation("02/01/2006, 15:04:05", match[1], location)
		if err != nil {
			return chat, nil, fmt.Errorf("parse export local date: %w", err)
		}

		messages = append(messages, WhatsAppMessage{
			Sender: match[2],
			Text:   match[3],
			Time:   date,
			Mine:   owner != "" && match[2] == owner && match[2] != chat.Name,
		})

		if match[2] == chat.Name {
			chat.Group = true
		}
	}

	if scanner.Err() != nil {
		return chat, nil, fmt.Errorf("read chat text: %w", scanner.Err())
	}

	if len(messages) == 0 {
		return chat, nil, errors.New("read export: no dated messages")
	}

	slices.SortStableFunc(messages, func(a, b WhatsAppMessage) int {
		return a.Time.Compare(b.Time)
	})

	return chat, messages, nil
}