package hazmat

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
)

type configuredAPIKeySource string

const (
	configuredAPIKeySourceNone   configuredAPIKeySource = ""
	configuredAPIKeySourceStore  configuredAPIKeySource = "host-secret-store"
	configuredAPIKeySourceLegacy configuredAPIKeySource = "legacy-agent-zshrc"
)

var agentZshrcPath = filepath.Join(agentHome, ".zshrc")

func secretStoreDirForHome(home string) string {
	return filepath.Join(home, ".hazmat", "secrets")
}

func providerSecretStorePath(envVar string) (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", fmt.Errorf("determine home directory for Hazmat secrets: %w", err)
	}
	return providerSecretStorePathForHome(home, envVar)
}

func providerSecretStorePathForHome(home, envVar string) (string, error) {
	descriptor, ok := providerCredentialDescriptorForEnvVar(envVar)
	if !ok {
		return "", fmt.Errorf("no host secret-store mapping for %s", envVar)
	}
	return providerSecretStorePathForDescriptor(home, descriptor)
}

func claudeCredentialStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessClaudeCredentials)
}

func claudeStateStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessClaudeState)
}

func codexAuthStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessCodexAuth)
}

func openCodeAuthStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessOpenCodeAuth)
}

func geminiOAuthStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessGeminiOAuth)
}

func geminiAccountsStorePathForHome(home string) string {
	return mustCredentialStorePathForHome(home, credentialHarnessGeminiAccounts)
}

func usesManagedAgentPath(path string) bool {
	return path == agentHome || isWithinDir(agentHome, path)
}

var agentPathForDirectIO = func(path string) string {
	return path
}

func readAgentSecretFile(path string) ([]byte, bool, error) {
	raw, err := os.ReadFile(agentPathForDirectIO(path))
	if err == nil {
		return raw, true, nil
	}
	if os.IsNotExist(err) {
		return nil, false, nil
	}
	if !usesManagedAgentPath(path) {
		return nil, false, err
	}

	out, agentErr := newAgentCommand("cat", path).Output()
	if agentErr != nil {
		if existsErr := newAgentCommand("test", "-e", path).Run(); agentTestReportsAbsent(existsErr) {
			return nil, false, nil
		}
		return nil, false, agentErr
	}
	return out, true, nil
}

func agentTestReportsAbsent(err error) bool {
	var exitErr *exec.ExitError
	return errors.As(err, &exitErr) && exitErr.ExitCode() == 1
}

func writeAgentSecretFile(path string, raw []byte, mode os.FileMode) error {
	if !usesManagedAgentPath(path) {
		if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
			return err
		}
		if err := os.WriteFile(path, raw, mode); err != nil {
			return err
		}
		return os.Chmod(path, mode)
	}
	if err := agentMkdirAll(filepath.Dir(path)); err != nil {
		return err
	}
	return agentWriteFile(path, raw, mode)
}

func removeAgentSecretFile(path string) error {
	if !usesManagedAgentPath(path) {
		if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
			return err
		}
		return nil
	}
	if err := asAgentQuiet("rm", "-f", path); err != nil {
		return err
	}
	return nil
}

func writeJSONMapStoreFile(path string, payload map[string]json.RawMessage) error {
	raw, err := json.MarshalIndent(payload, "", "  ")
	if err != nil {
		return err
	}
	raw = append(raw, '\n')
	return writeHostStoredSecretFile(path, raw)
}

func readJSONMapStoreFile(path string) (map[string]json.RawMessage, bool, error) {
	raw, ok, err := readHostStoredSecretFile(path)
	if err != nil || !ok {
		return nil, ok, err
	}
	var payload map[string]json.RawMessage
	if err := json.Unmarshal(raw, &payload); err != nil {
		return nil, false, err
	}
	if len(payload) == 0 {
		return nil, false, nil
	}
	return payload, true, nil
}

func readClaudeStateKeysFromAgent(path string) (map[string]json.RawMessage, bool, error) {
	raw, ok, err := readAgentSecretFile(path)
	if err != nil || !ok {
		return nil, ok, err
	}
	payload, err := selectClaudeAuthKeys(raw)
	if err != nil {
		return nil, false, err
	}
	if len(payload) == 0 {
		return nil, false, nil
	}
	return payload, true, nil
}

func writeClaudeStateKeysToAgent(path string, payload map[string]json.RawMessage) error {
	currentRaw, ok, err := readAgentSecretFile(path)
	if err != nil {
		return err
	}

	current := map[string]json.RawMessage{}
	if ok && len(bytes.TrimSpace(currentRaw)) > 0 {
		if err := json.Unmarshal(currentRaw, &current); err != nil {
			return err
		}
	}
	for key, value := range payload {
		current[key] = value
	}

	raw, err := json.MarshalIndent(current, "", "  ")
	if err != nil {
		return err
	}
	raw = append(raw, '\n')
	return writeAgentSecretFile(path, raw, 0o600)
}

func removeClaudeStateKeysFromAgent(path string) error {
	raw, ok, err := readAgentSecretFile(path)
	if err != nil || !ok {
		return err
	}

	var current map[string]json.RawMessage
	if err := json.Unmarshal(raw, &current); err != nil {
		return err
	}
	for _, key := range claudePortableAuthKeys {
		delete(current, key)
	}
	if len(current) == 0 {
		return removeAgentSecretFile(path)
	}

	updated, err := json.MarshalIndent(current, "", "  ")
	if err != nil {
		return err
	}
	updated = append(updated, '\n')
	return writeAgentSecretFile(path, updated, 0o600)
}

func readHostStoredAPIKey(spec harnessAPIKeySpec) (string, error) {
	path, err := providerSecretStorePath(spec.EnvVar)
	if err != nil {
		return "", err
	}

	raw, err := os.ReadFile(path)
	if err != nil {
		if os.IsNotExist(err) {
			return "", nil
		}
		return "", fmt.Errorf("read %s from host secret store: %w", spec.EnvVar, err)
	}
	return strings.TrimSpace(string(raw)), nil
}

func storeHostAPIKey(spec harnessAPIKeySpec, value string) error {
	path, err := providerSecretStorePath(spec.EnvVar)
	if err != nil {
		return err
	}

	value = strings.TrimSpace(value)
	if value == "" {
		return fmt.Errorf("%s cannot be empty", spec.EnvVar)
	}

	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		return fmt.Errorf("create host secret-store directory for %s: %w", spec.EnvVar, err)
	}
	if err := os.WriteFile(path, []byte(value+"\n"), 0o600); err != nil {
		return fmt.Errorf("write %s to host secret store: %w", spec.EnvVar, err)
	}
	return nil
}

func lookupConfiguredAPIKey(spec harnessAPIKeySpec) (string, configuredAPIKeySource, error) {
	value, err := readHostStoredAPIKey(spec)
	if err != nil {
		return "", configuredAPIKeySourceNone, err
	}
	if value != "" {
		return value, configuredAPIKeySourceStore, nil
	}

	line := readZshrcEnvLine(agentZshrcPath, spec.EnvVar)
	if line == "" {
		return "", configuredAPIKeySourceNone, nil
	}

	value, ok := parseExportedEnvLineValue(line, spec.EnvVar)
	if !ok || strings.TrimSpace(value) == "" {
		return "", configuredAPIKeySourceNone, nil
	}
	return value, configuredAPIKeySourceLegacy, nil
}

func parseExportedEnvLineValue(line, envVar string) (string, bool) {
	trimmed := strings.TrimSpace(line)
	prefix := "export " + envVar + "="
	if !strings.HasPrefix(trimmed, prefix) {
		return "", false
	}

	raw := strings.TrimSpace(strings.TrimPrefix(trimmed, prefix))
	if raw == "" {
		return "", false
	}
	if value, err := strconv.Unquote(raw); err == nil {
		return value, true
	}
	return strings.Trim(raw, "\"'"), true
}

func harnessAPIKeyPromptByEnvVar(envVar string) harnessAPIKeySpec {
	for _, spec := range harnessAPIKeyPrompts {
		if spec.EnvVar == envVar {
			return spec
		}
	}
	return harnessAPIKeySpec{EnvVar: envVar}
}

func removableLegacyAPIKeyEnvVars(pending []pendingAPIKeyUpdate) []string {
	seen := make(map[string]struct{}, len(harnessAPIKeyPrompts)+len(pending))
	var envVars []string

	for _, update := range pending {
		if _, dup := seen[update.EnvVar]; dup {
			continue
		}
		seen[update.EnvVar] = struct{}{}
		envVars = append(envVars, update.EnvVar)
	}

	for _, spec := range harnessAPIKeyPrompts {
		if _, dup := seen[spec.EnvVar]; dup {
			continue
		}
		value, err := readHostStoredAPIKey(spec)
		if err != nil || value == "" {
			continue
		}
		seen[spec.EnvVar] = struct{}{}
		envVars = append(envVars, spec.EnvVar)
	}

	return envVars
}

func apiKeyRemovalUpdates(envVars []string) []pendingAPIKeyUpdate {
	updates := make([]pendingAPIKeyUpdate, 0, len(envVars))
	for _, envVar := range envVars {
		updates = append(updates, pendingAPIKeyUpdate{EnvVar: envVar})
	}
	return updates
}

func hasLegacyAPIKeyExports(envVars []string) bool {
	for _, envVar := range envVars {
		if readZshrcEnvLine(agentZshrcPath, envVar) != "" {
			return true
		}
	}
	return false
}

func removeLegacyAPIKeyExports(envVars []string) error {
	return updateAgentFile(
		agentZshrcPath,
		func(content string) string {
			return rewriteZshrcAPIKeys(content, apiKeyRemovalUpdates(envVars))
		},
		0o600,
	)
}

func applyHarnessAPIKeyEnv(cfg *sessionConfig) error {
	return applyHarnessAPIKeyEnvForSession(cfg, false)
}

func applyHarnessAPIKeyEnvForSession(cfg *sessionConfig, planOnly bool) error {
	descriptors := providerCredentialDescriptorsForHarness(cfg.HarnessID)
	for _, descriptor := range descriptors {
		if err := applyProviderAPIKeyEnvForSession(cfg, descriptor, planOnly); err != nil {
			return err
		}
	}
	return nil
}

func applyProviderAPIKeyEnvForSession(cfg *sessionConfig, descriptor credentialDescriptor, planOnly bool) error {
	spec := harnessAPIKeyPromptByEnvVar(descriptor.EnvVar)
	value, source, err := lookupConfiguredAPIKey(spec)
	if err != nil {
		return err
	}
	if value == "" {
		return nil
	}

	switch source {
	case configuredAPIKeySourceNone:
		return nil
	case configuredAPIKeySourceStore:
	case configuredAPIKeySourceLegacy:
		if planOnly {
			cfg.SessionNotes = append(cfg.SessionNotes,
				fmt.Sprintf("Legacy %s from %s would be migrated into ~/.hazmat/secrets before launch.", spec.EnvVar, agentZshrcPath))
			break
		}
		if err := storeHostAPIKey(spec, value); err != nil {
			cfg.SessionNotes = append(cfg.SessionNotes,
				fmt.Sprintf("Using legacy %s from %s because migration into ~/.hazmat/secrets failed: %v", spec.EnvVar, agentZshrcPath, err))
			break
		}
		if err := removeLegacyAPIKeyExports([]string{spec.EnvVar}); err != nil {
			cfg.SessionNotes = append(cfg.SessionNotes,
				fmt.Sprintf("Migrated legacy %s into ~/.hazmat/secrets, but could not remove the old export from %s: %v", spec.EnvVar, agentZshrcPath, err))
		} else {
			cfg.SessionNotes = append(cfg.SessionNotes,
				fmt.Sprintf("Migrated legacy %s from %s into ~/.hazmat/secrets.", spec.EnvVar, agentZshrcPath))
		}
	default:
		return nil
	}

	if cfg.HarnessEnv == nil {
		cfg.HarnessEnv = make(map[string]string, 1)
	}
	cfg.HarnessEnv[spec.EnvVar] = value
	grantSource := "host secret store"
	if planOnly && source == configuredAPIKeySourceLegacy {
		grantSource = "legacy agent zshrc"
	}
	cfg.CredentialEnvGrants = appendSessionCredentialEnvGrant(cfg.CredentialEnvGrants, sessionCredentialEnvGrant{
		EnvVar:          spec.EnvVar,
		CredentialID:    descriptor.ID,
		Source:          grantSource,
		ConsumerHarness: cfg.HarnessID,
	})
	return nil
}
