package model

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"sync/atomic"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/config/runtimeconfig"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// ModelRouter is the routing contract the agent (and loop) depend on. *Router
// implements it; RuntimeLLM implements it too, so a hot-swappable provider can
// serve as both the adapter and the router without a separate router reference
// to keep in sync.
type ModelRouter interface {
	RouteRequest(req Request) (model string, needsVision bool)
}

// runtimeChainSnapshot is an immutable, self-contained LLM stack snapshot. It is
// stored and swapped atomically: a request that already captured a snapshot
// keeps using that exact chain+router for its whole call, so a reload never
// interrupts an in-flight request.
type runtimeChainSnapshot struct {
	chain  *Chain
	router *Router
	group  runtimeconfig.Group
	groups []runtimeconfig.Group
	// providerChains holds one chain per *other* protocol declared by the active
	// group's personaProviders (裁决 03.4.1 的 per-persona provider：跨协议换
	// adapter). A persona whose provider binding resolves to one of these keys is
	// served by that chain; every other request keeps the default chain above.
	providerChains map[string]*Chain
}

// RuntimeLLM is the hot-swappable LLM provider stack for the runtime. It
// implements Adapter (delegating to the current chain) and ModelRouter
// (delegating to the current router). A panel save reloads models.json, rebuilds
// a fresh chain+router and atomically swaps it in: new calls use the new stack
// while calls already in flight finish on the old one.
type RuntimeLLM struct {
	path    string
	env     config.LLMConfig
	log     *slog.Logger
	current atomic.Value // *runtimeChainSnapshot
}

// NewRuntimeLLM builds the provider stack from models.json at path, falling back
// to env. It loads the initial snapshot so the runtime is usable immediately.
// log may be nil (a discard logger is used).
func NewRuntimeLLM(path string, env config.LLMConfig, log *slog.Logger) *RuntimeLLM {
	if log == nil {
		log = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	m := &RuntimeLLM{path: path, env: env, log: logger.WithSource(log, "model")}
	m.loadInitial()
	return m
}

// Path returns the models.json path this manager reads and reloads from.
func (m *RuntimeLLM) Path() string { return m.path }

// Groups returns a copy of the current loaded groups; the active group (the one
// a new request uses by default) is index 0.
func (m *RuntimeLLM) Groups() []runtimeconfig.Group {
	src := m.load().groups
	out := make([]runtimeconfig.Group, len(src))
	copy(out, src)
	return out
}

// Group returns the active group (index 0).
func (m *RuntimeLLM) Group() runtimeconfig.Group { return m.load().group }

// Current returns the current chain and router. The returned chain stays valid
// for calls already in flight; new calls pick up the latest swap.
func (m *RuntimeLLM) Current() (*Chain, *Router) {
	snap := m.load()
	return snap.chain, snap.router
}

// Reload re-reads models.json (falling back to env) and atomically rebuilds the
// in-use chain+router. In-flight calls keep the previous snapshot.
func (m *RuntimeLLM) Reload() error {
	groups, warn := runtimeconfig.LoadModels(m.path, m.env)
	if warn != nil {
		m.log.Warn("models.json reload fell back to env", "error", warn)
	}
	return m.swap(groups)
}

// Provider returns the active group's provider protocol.
func (m *RuntimeLLM) Provider() string { return m.load().group.Provider }

// Model returns the active group's primary model name.
func (m *RuntimeLLM) Model() string { return m.load().group.Model }

// Name implements Adapter, returning the active group id.
func (m *RuntimeLLM) Name() string { return m.load().group.ID }

// Chat implements Adapter, delegating to the chain the request's persona
// resolves to. A request that began before a reload completes on the snapshot it
// captured. When no model is configured (nil chain) it returns ErrNoModel rather
// than panicking.
func (m *RuntimeLLM) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
	chain := m.chainFor(req.PersonaID)
	if chain == nil {
		return Response{}, ErrNoModel
	}
	return chain.Chat(ctx, req, opts)
}

// ChatStream implements Adapter.
func (m *RuntimeLLM) ChatStream(ctx context.Context, req Request, onEvent func(StreamEvent)) (Response, error) {
	chain := m.chainFor(req.PersonaID)
	if chain == nil {
		return Response{}, ErrNoModel
	}
	return chain.ChatStream(ctx, req, onEvent)
}

// chainFor resolves the chain that serves this persona (裁决 03.4.1 的
// per-persona provider). It consumes Router.RouteProviderForPersona — the
// provider decision the router already made for this persona — and only switches
// when BOTH hold:
//
//   - the persona has its own model binding (an entry in personaProviders alone
//     would send the group's default model name to another vendor's endpoint),
//   - and a chain exists for that provider (a declared protocol with no matching
//     group in models.json is reported when the snapshot is built, and the
//     persona then keeps the default chain rather than failing the turn).
//
// Everything else — no persona, no binding, same protocol as the active group —
// returns the default chain, so a deployment without per-persona providers
// behaves exactly as before.
func (m *RuntimeLLM) chainFor(persona string) *Chain {
	snap := m.load()
	if p := trimKey(persona); p != "" && snap.router != nil {
		if _, _, ok := snap.group.PersonaRoute(p); ok {
			if c := snap.providerChains[snap.router.RouteProviderForPersona(p)]; c != nil {
				return c
			}
		}
	}
	return snap.chain
}

// RouteRequest implements ModelRouter, delegating to the current router.
func (m *RuntimeLLM) RouteRequest(req Request) (string, bool) {
	r := m.load().router
	if r == nil {
		return "", false
	}
	return r.RouteRequest(req)
}

// RouteRequestForPersona implements PersonaRouter (裁决 03.4.1): the loop routes
// a turn as the persona it answers as, and a per-persona binding in the active
// models.json group reaches the request through here.
func (m *RuntimeLLM) RouteRequestForPersona(personaID string, req Request) (string, bool) {
	r := m.load().router
	if r == nil {
		return "", false
	}
	return r.RouteRequestForPersona(personaID, req)
}

// loadInitial builds the first snapshot; a build failure leaves a nil chain so
// calls report a clear no-model error instead of panicking.
func (m *RuntimeLLM) loadInitial() {
	groups, warn := runtimeconfig.LoadModels(m.path, m.env)
	if warn != nil {
		m.log.Warn("models.json load fell back to env", "error", warn)
	}
	_ = m.swap(groups)
}

// swap rebuilds and atomically stores a snapshot for groups. A blank slice falls
// back to the env-synthesized group so there is always at least one.
func (m *RuntimeLLM) swap(groups []runtimeconfig.Group) error {
	if len(groups) == 0 {
		groups = runtimeconfig.GroupFromEnv(m.env)
	}
	snap, err := buildSnapshot(groups, m.env, m.log)
	if err != nil {
		return err
	}
	m.current.Store(snap)
	return nil
}

// load returns the current snapshot, or an empty (nil chain) one when the
// manager was never initialized.
func (m *RuntimeLLM) load() *runtimeChainSnapshot {
	if v := m.current.Load(); v != nil {
		return v.(*runtimeChainSnapshot)
	}
	return &runtimeChainSnapshot{}
}

// buildSnapshot turns the active group (index 0) into a chain+router. A group
// without a base URL/model yields a nil chain (no model configured), surfaced as
// ErrNoModel by the adapter methods.
//
// The env-local rung is copied onto the group's flat config because a models.json
// group does not carry it: without that copy the third degradation level
// (裁决 03.4.3) would silently disappear as soon as models.json exists.
//
// The router is seeded with the active group's per-persona bindings
// (裁决 03.4.1). It is deliberately NOT given the env vision model: image routing
// is owned by vision.Router, and a "has image → switch model" rule here bypassed
// that decision (裁决 14.4.1/14.4.5).
func buildSnapshot(groups []runtimeconfig.Group, env config.LLMConfig, log *slog.Logger) (*runtimeChainSnapshot, error) {
	active := groups[0]
	cfg := active.ToLLMConfig()
	cfg.Local = env.Local
	chain, err := NewChatChainWithLog(cfg, log)
	if err != nil && !errors.Is(err, ErrNoModel) {
		return nil, err
	}
	if chain != nil {
		// Budgets are resolved inside the chain: the active group's own
		// maxTokens (models.json) wins, and env.FallbackMaxTokens only caps the
		// last step of a chain that really has a fallback (design/04 §3.1).
		chain = chain.WithTokenBudget(active.MaxTokens, env.FallbackMaxTokens)
	}
	router := NewRouter(DefaultModelForProvider(cfg), personaRouterOptions(active, cfg)...)
	return &runtimeChainSnapshot{
		chain:          chain,
		router:         router,
		group:          active,
		groups:         groups,
		providerChains: personaProviderChains(groups, active, env, log),
	}, nil
}

// personaProviderChains builds one chain per protocol that the active group's
// personaProviders declares and that models.json actually provides a group for
// (裁决 03.4.1 的 per-persona provider：跨协议换 adapter). The persona's own
// model name is pinned as the chain's primary model, so a caller that stamps the
// persona without a routed model name still reaches the right endpoint.
//
// A persona that names the same protocol as the active group needs no chain of
// its own: the default chain already speaks that protocol, and the per-persona
// model name reaches it through the router. A personaProviders entry whose
// protocol has no group is reported and skipped — a half-built chain could only
// fail, and the persona keeps the default chain instead.
func personaProviderChains(groups []runtimeconfig.Group, active runtimeconfig.Group, env config.LLMConfig, log *slog.Logger) map[string]*Chain {
	if len(active.PersonaProviders) == 0 {
		return nil
	}
	out := make(map[string]*Chain, len(active.PersonaProviders))
	for persona := range active.PersonaProviders {
		model, provider, ok := active.PersonaRoute(persona)
		if !ok {
			warn(log, "personaProviders entry ignored: the persona binds no model of its own",
				"persona", persona, "field", "personaModels")
			continue
		}
		if provider == active.Provider {
			continue
		}
		if _, built := out[provider]; built {
			continue
		}
		g, found := groupForProvider(groups, provider)
		if !found {
			warn(log, "per-persona provider skipped: no models.json group speaks it",
				"persona", persona, "provider", provider)
			continue
		}
		chain := personaChain(g, model, env, log)
		if chain != nil {
			out[provider] = chain
		}
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

// personaChain builds and budgets the chain serving one persona through another
// vendor's group: the group's flat config with the persona's own model name as
// the primary model, plus the shared local last rung. The reasoning shape
// (reasoning.kind) travels with the group's config, so a persona routed to
// another vendor keeps that vendor's switch shape. It reports the reason and
// returns nil when the group cannot serve a chain.
func personaChain(g runtimeconfig.Group, model string, env config.LLMConfig, log *slog.Logger) *Chain {
	cfg := g.ToLLMConfig()
	if cfg.Provider == ProviderAnthropic {
		cfg.AnthropicModel = model
	} else {
		cfg.Model = model
	}
	cfg.Local = env.Local
	chain, err := NewChatChainWithLog(cfg, log)
	if err != nil && !errors.Is(err, ErrNoModel) {
		warn(log, "per-persona chain skipped", "group", g.ID, "model", model, "error", err)
		return nil
	}
	if chain == nil {
		warn(log, "per-persona chain skipped: the provider group has no base url or model",
			"group", g.ID, "model", model)
		return nil
	}
	return chain.WithTokenBudget(g.MaxTokens, env.FallbackMaxTokens)
}

// groupForProvider returns the first group speaking the given protocol. The group
// order in models.json is the operator's choice, so "first match" is a stable,
// documentable rule rather than a silent pick — a duplicate protocol is reported
// by the builder when it matters.
func groupForProvider(groups []runtimeconfig.Group, provider string) (runtimeconfig.Group, bool) {
	for _, g := range groups {
		if g.Provider == provider {
			return g, true
		}
	}
	return runtimeconfig.Group{}, false
}

// personaRouterOptions turns the active group's per-persona model/provider
// bindings into router options (裁决 03.4.1). It is a helper so buildSnapshot
// stays inside the 80-line function budget.
func personaRouterOptions(active runtimeconfig.Group, cfg config.LLMConfig) []RouterOption {
	opts := []RouterOption{WithDefaultProvider(cfg.Provider)}
	for persona, model := range active.PersonaModels {
		opts = append(opts, WithPersonaModel(persona, model))
	}
	for persona, provider := range active.PersonaProviders {
		opts = append(opts, WithPersonaProvider(persona, provider))
	}
	return opts
}