// plexum-gemini-provider-plugin is a Plexum ProviderPlugin that wraps
// the local `gemini` CLI binary (Google Gemini CLI). Each Plexum turn
// forks one `gemini -p` subprocess in the agent's workspace, captures
// JSON output, and emits canonical TurnEvents. Multi-turn continuity
// via `--resume <session-id>` (session id stashed at
// <workspace>/.plexum-gemini-session).
//
// Reference: openclaw/extensions/google/cli-backend.ts. Same CLI-driven
// approach as Plexum-anthropic-provider's contractor mode.
//
// Models — Plexum agent.json.model can be any of:
//   gemini-pro          → CLI alias "pro"          → gemini-3.1-pro-preview
//   gemini-flash        → CLI alias "flash"        → gemini-3.1-flash-preview (default)
//   gemini-flash-lite   → CLI alias "flash-lite"   → gemini-3.1-flash-lite-preview
//   gemini              → no --model flag (gemini-cli default)
// or any raw model id the local gemini CLI accepts (passed through).
package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"

	"git.hangman-lab.top/hzhang/Plexum-sdk-go/canonical"
	"git.hangman-lab.top/hzhang/Plexum-sdk-go/mcpbridge"
	plugin "git.hangman-lab.top/hzhang/Plexum-sdk-go/plugin"

	"git.hangman-lab.top/hzhang/Plexum-gemini-provider/internal/runner"
)

const pluginName = "plexum-gemini-provider"

// Aliases the plugin translates before passing --model to the CLI.
// CLI accepts the aliases directly too, but normalizing keeps logs
// consistent across operator + provider sides.
var modelAliases = map[string]string{
	"gemini-pro":        "pro",
	"gemini-flash":      "flash",
	"gemini-flash-lite": "flash-lite",
}

// supportedModels: what the manifest's provider.models advertises.
// Operator agent.json.model must match one of these; we route the
// "gemini" plain id to the CLI's default (no --model flag).
var supportedModels = []string{
	"gemini",
	"gemini-pro",
	"gemini-flash",
	"gemini-flash-lite",
}

// HostConfig is per-profile plugin config at
// <profile>/plugins/plexum-gemini-provider/config.json. All fields
// optional — sensible defaults if file is missing.
//
//	{
//	  "binary":     "gemini",         // executable name or absolute path
//	  "extra_args": ["--yolo"]        // appended after the standard flags
//	}
//
// No api_key field: gemini CLI auth lives in ~/.gemini (OAuth or
// GEMINI_API_KEY env), already set up by the operator.
type HostConfig struct {
	Binary    string   `json:"binary"`
	ExtraArgs []string `json:"extra_args"`
}

type geminiPlugin struct {
	host plugin.HostAPI
	cfg  HostConfig
}

func (p *geminiPlugin) Manifest() plugin.Manifest {
	return plugin.Manifest{
		Name:       pluginName,
		Version:    "0.1.0",
		Activation: plugin.ActivationLazy,
		Executable: "plexum-gemini-provider-plugin",
		Contracts: plugin.Contracts{
			Provider: &plugin.ProviderContract{Models: supportedModels},
		},
	}
}

func (p *geminiPlugin) Init(ctx context.Context, host plugin.HostAPI) error {
	p.host = host

	profileRoot := os.Getenv("PLEXUM_PROFILE_ROOT")
	if profileRoot == "" {
		home, _ := os.UserHomeDir()
		profileRoot = filepath.Join(home, ".plexum")
	}
	cfgPath := filepath.Join(profileRoot, "plugins", pluginName, "config.json")
	raw, err := os.ReadFile(cfgPath)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return fmt.Errorf("read %s: %w", cfgPath, err)
	}
	if len(raw) > 0 {
		if err := json.Unmarshal(raw, &p.cfg); err != nil {
			return fmt.Errorf("parse %s: %w", cfgPath, err)
		}
	}
	if p.cfg.Binary == "" {
		p.cfg.Binary = "gemini"
	}
	host.Log("info", "gemini provider initialized", map[string]any{
		"binary": p.cfg.Binary, "extra_args": p.cfg.ExtraArgs, "models": supportedModels,
	})
	return nil
}

// Stream is the bare ProviderPlugin entrypoint; called when the SDK
// can't see StreamWithAgent. We use the agent-aware variant because
// the CLI needs the workspace cwd.
func (p *geminiPlugin) Stream(ctx context.Context, modelID string, req canonical.TurnRequest) (<-chan canonical.TurnEvent, error) {
	return p.StreamWithAgent(ctx, modelID, plugin.AgentContext{}, req)
}

// StreamWithAgent receives the AgentContext (workspace path) Plexum's
// host wires in for this kind of plugin (see ProviderPluginWithAgent).
func (p *geminiPlugin) StreamWithAgent(ctx context.Context, modelID string, agent plugin.AgentContext, req canonical.TurnRequest) (<-chan canonical.TurnEvent, error) {
	cliModel := mapModel(modelID)
	return runner.Run(ctx, p.host, agent, cliModel, p.cfg.Binary, p.cfg.ExtraArgs, req)
}

// MutateSession mirrors host dynamic-* consume mutations onto the
// gemini chat JSONL at ~/.gemini/tmp/<ws>/chats/session-*.jsonl so
// the next `gemini --resume` sees the consumed view. Tool-call
// round-trip is captured by runner.ParseGeminiToolCalls +
// EmitGeminiToolCalls (post-turn), so the BlockMutations we receive
// here resolve to ids actually in gemini's session.
func (p *geminiPlugin) MutateSession(ctx context.Context, req plugin.SessionMutateRequest) error {
	if len(req.Mutations) == 0 {
		return nil
	}
	mu := make([]runner.GeminiMutation, 0, len(req.Mutations))
	for _, m := range req.Mutations {
		mu = append(mu, runner.GeminiMutation{BlockID: m.BlockID, Op: m.Op})
	}
	touched, err := runner.MutateGeminiSession(req.Workspace, mu)
	if err != nil {
		p.host.Log("warn", "gemini session mutate failed", map[string]any{
			"agent": req.AgentID, "err": err.Error(),
		})
		return err
	}
	p.host.Log("info", "gemini session mutate complete", map[string]any{
		"agent": req.AgentID, "touched": touched, "requested": len(req.Mutations),
	})
	return nil
}

// mapModel converts a Plexum-advertised model id into the CLI's
// --model value. Returns "" for "gemini" (use CLI default) and
// passes any other unrecognized id through verbatim (operator may
// have configured a specific gemini-3.x model id directly).
func mapModel(plexumModel string) string {
	if plexumModel == "" || plexumModel == "gemini" {
		return ""
	}
	if alias, ok := modelAliases[plexumModel]; ok {
		return alias
	}
	return plexumModel
}

func main() {
	if len(os.Args) > 1 && os.Args[1] == "mcp-host" {
		if err := runMCPHost(os.Args[2:]); err != nil {
			fmt.Fprintf(os.Stderr, "mcp-host: %v\n", err)
			os.Exit(1)
		}
		return
	}
	if err := plugin.Serve(&geminiPlugin{}); err != nil {
		fmt.Fprintf(os.Stderr, "plexum-gemini-provider-plugin: %v\n", err)
		os.Exit(1)
	}
}

func runMCPHost(argv []string) error {
	opts := mcpbridge.Opts{
		SocketPath: os.Getenv("PLEXUM_MCP_SOCKET"),
		AgentID:    os.Getenv("PLEXUM_MCP_AGENT_ID"),
	}
	for i := 0; i < len(argv); i++ {
		switch argv[i] {
		case "--socket":
			if i+1 >= len(argv) {
				return errors.New("--socket needs value")
			}
			opts.SocketPath = argv[i+1]
			i++
		case "--agent-id":
			if i+1 >= len(argv) {
				return errors.New("--agent-id needs value")
			}
			opts.AgentID = argv[i+1]
			i++
		default:
			return fmt.Errorf("unknown mcp-host flag: %s", argv[i])
		}
	}
	return mcpbridge.Run(context.Background(), opts)
}