package repositorylocalproxy

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"encoding/base64"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"path"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	CapabilityRepoRead         = "repo:read"
	CapabilityRepoClone        = "repo:clone"
	CapabilityRepoWrite        = "repo:write"
	CapabilityPullRequestRead  = "pull_request:read"
	CapabilityPullRequestWrite = "pull_request:write"

	EnvLeaseToken   = "AGENTTEAM_REPOSITORY_ACCESS_LOCAL_PROXY_TOKEN"
	EnvProxyBase    = "AGENTTEAM_REPOSITORY_ACCESS_LOCAL_PROXY_BASE_URL"
	EnvProviderHost = "AGENTTEAM_REPOSITORY_ACCESS_PROVIDER_HOST"
)

var secretLike = regexp.MustCompile(`(?i)(agentteam_lp_[A-Za-z0-9_-]+|ags_[A-Za-z0-9_-]+|gh[pousr]_[A-Za-z0-9_]+|github_pat_[A-Za-z0-9_]+)`)

type RepositoryScope struct {
	ProviderHost string
	Repositories []string
}

type LeaseRequest struct {
	MachineID    string
	AgentID      string
	SessionID    string
	WorkspaceID  string
	Scope        RepositoryScope
	Capabilities []string
	AGSTokenID   string
	TTL          time.Duration
	Now          time.Time
}

type Lease struct {
	ID           string
	Token        string
	MachineID    string
	AgentID      string
	SessionID    string
	WorkspaceID  string
	Scope        RepositoryScope
	Capabilities []string
	AGSTokenID   string
	IssuedAt     time.Time
	ExpiresAt    time.Time
	RevokedAt    time.Time
}

type LeaseStore struct {
	mu         sync.RWMutex
	byToken    map[string]Lease
	byID       map[string]string
	tokenBytes func() ([]byte, error)
}

func NewLeaseStore() *LeaseStore {
	return &LeaseStore{byToken: map[string]Lease{}, byID: map[string]string{}, tokenBytes: randomTokenBytes}
}

func (s *LeaseStore) Issue(req LeaseRequest) (Lease, error) {
	if s == nil {
		return Lease{}, errors.New("lease store is unavailable")
	}
	now := req.Now
	if now.IsZero() {
		now = time.Now().UTC()
	}
	ttl := req.TTL
	if ttl <= 0 {
		ttl = 30 * time.Minute
	}
	if strings.TrimSpace(req.AgentID) == "" || strings.TrimSpace(req.SessionID) == "" || strings.TrimSpace(req.WorkspaceID) == "" {
		return Lease{}, errors.New("lease identity is incomplete")
	}
	host := normalizeHost(req.Scope.ProviderHost)
	if host == "" {
		return Lease{}, errors.New("lease provider host is required")
	}
	repos := normalizeRepos(req.Scope.Repositories)
	if len(repos) == 0 {
		return Lease{}, errors.New("lease repository scope is required")
	}
	caps := normalizeCapabilities(req.Capabilities)
	if len(caps) == 0 {
		return Lease{}, errors.New("lease capabilities are required")
	}
	b, err := s.tokenBytes()
	if err != nil {
		return Lease{}, err
	}
	tok := "agentteam_lp_" + base64.RawURLEncoding.EncodeToString(b)
	lease := Lease{ID: "lpl_" + base64.RawURLEncoding.EncodeToString(b[:min(12, len(b))]), Token: tok, MachineID: strings.TrimSpace(req.MachineID), AgentID: strings.TrimSpace(req.AgentID), SessionID: strings.TrimSpace(req.SessionID), WorkspaceID: strings.TrimSpace(req.WorkspaceID), Scope: RepositoryScope{ProviderHost: host, Repositories: repos}, Capabilities: caps, AGSTokenID: strings.TrimSpace(req.AGSTokenID), IssuedAt: now, ExpiresAt: now.Add(ttl)}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.byToken[tok] = lease
	s.byID[lease.ID] = tok
	return lease, nil
}

func (s *LeaseStore) Revoke(id string, now time.Time) bool {
	if s == nil {
		return false
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	tok, ok := s.byID[strings.TrimSpace(id)]
	if !ok {
		tok = strings.TrimSpace(id)
	}
	l, ok := s.byToken[tok]
	if !ok {
		return false
	}
	l.RevokedAt = now
	s.byToken[tok] = l
	return true
}

func (s *LeaseStore) Validate(token string, now time.Time) (Lease, error) {
	if s == nil {
		return Lease{}, ErrInvalidLease
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	token = strings.TrimSpace(token)
	if token == "" {
		return Lease{}, ErrInvalidLease
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	for key, l := range s.byToken {
		if subtle.ConstantTimeCompare([]byte(key), []byte(token)) == 1 {
			if !l.RevokedAt.IsZero() {
				return Lease{}, ErrRevokedLease
			}
			if !l.ExpiresAt.IsZero() && !now.Before(l.ExpiresAt) {
				return Lease{}, ErrExpiredLease
			}
			return l, nil
		}
	}
	return Lease{}, ErrInvalidLease
}

var ErrInvalidLease = errors.New("invalid local repository access lease")
var ErrExpiredLease = errors.New("expired local repository access lease")
var ErrRevokedLease = errors.New("revoked local repository access lease")
var ErrForbidden = errors.New("repository access forbidden")
var ErrUnknownMapping = errors.New("unknown repository mapping")
var ErrUpstreamToken = errors.New("upstream repository token unavailable")

type UpstreamCredentialResolver interface {
	ResolveRepositoryCredential(context.Context, Lease, string, string, string) (UpstreamCredential, error)
}
type UpstreamCredential struct {
	Token   string
	Revoked bool
	Expired bool
}

type AuditEvent struct {
	LeaseID    string
	AgentID    string
	Host       string
	Repository string
	Capability string
	Decision   string
	Reason     string
	At         time.Time
}

type AuditRecorder interface {
	RecordRepositoryAccessProxyEvent(context.Context, AuditEvent)
}

type Proxy struct {
	Store           *LeaseStore
	Resolver        UpstreamCredentialResolver
	UpstreamBaseURL string
	Now             func() time.Time
	Logger          *slog.Logger
	Client          *http.Client
	Audit           AuditRecorder
}

func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	now := time.Now().UTC()
	if p.Now != nil {
		now = p.Now()
	}
	lease, err := p.Store.Validate(extractLeaseToken(r), now)
	if err != nil {
		p.record(r.Context(), AuditEvent{Decision: "denied", Reason: err.Error(), At: now})
		p.deny(w, http.StatusUnauthorized, err, "lease rejected")
		return
	}
	host, repo, capability, upPath, err := mapRequest(r, lease)
	if err != nil {
		p.record(r.Context(), AuditEvent{LeaseID: lease.ID, AgentID: lease.AgentID, Decision: "denied", Reason: err.Error(), At: now})
		p.deny(w, statusFor(err), err, "request mapping rejected")
		return
	}
	cred, err := p.Resolver.ResolveRepositoryCredential(r.Context(), lease, host, repo, capability)
	if err != nil || cred.Revoked || cred.Expired || strings.TrimSpace(cred.Token) == "" {
		p.Store.Revoke(lease.ID, now)
		p.record(r.Context(), AuditEvent{LeaseID: lease.ID, AgentID: lease.AgentID, Host: host, Repository: repo, Capability: capability, Decision: "denied", Reason: ErrUpstreamToken.Error(), At: now})
		p.deny(w, http.StatusForbidden, ErrUpstreamToken, "upstream token rejected")
		return
	}
	base, err := url.Parse(strings.TrimRight(p.UpstreamBaseURL, "/"))
	if err != nil || base.Scheme == "" || base.Host == "" {
		p.deny(w, http.StatusBadGateway, errors.New("upstream repository proxy is not configured"), "upstream unavailable")
		return
	}
	outURL := *base
	outURL.Path = joinURLPath(base.Path, upPath)
	outURL.RawQuery = sanitizedUpstreamRawQuery(r.URL.Query())
	outReq, err := http.NewRequestWithContext(r.Context(), r.Method, outURL.String(), r.Body)
	if err != nil {
		p.deny(w, http.StatusBadGateway, err, "upstream request failed")
		return
	}
	copyHeaders(outReq.Header, r.Header)
	outReq.Header.Set("Authorization", "Bearer "+cred.Token)
	client := p.Client
	if client == nil {
		client = http.DefaultClient
	}
	p.record(r.Context(), AuditEvent{LeaseID: lease.ID, AgentID: lease.AgentID, Host: host, Repository: repo, Capability: capability, Decision: "forwarded", At: now})
	resp, err := client.Do(outReq)
	if err != nil {
		p.deny(w, http.StatusBadGateway, err, "upstream request failed")
		return
	}
	defer resp.Body.Close()
	copyHeaders(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)
	_, _ = io.Copy(w, resp.Body)
}

func (p *Proxy) record(ctx context.Context, event AuditEvent) {
	if p != nil && p.Audit != nil {
		p.Audit.RecordRepositoryAccessProxyEvent(ctx, event)
	}
}

func (p *Proxy) deny(w http.ResponseWriter, status int, err error, msg string) {
	if p.Logger != nil {
		p.Logger.Warn(Redact(msg), "error", Redact(err.Error()))
	}
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(map[string]string{"error": Redact(err.Error())})
}

func mapRequest(r *http.Request, lease Lease) (host, repo, capability, upstreamPath string, err error) {
	parts := splitPath(r.URL.Path)
	if len(parts) >= 4 && parts[0] == "_agentteam_git" {
		host = normalizeHost(parts[1])
		repo = normalizeRepo(parts[2] + "/" + strings.TrimSuffix(parts[3], ".git"))
		upstreamPath = "/" + strings.Join(parts[2:], "/")
		capability = CapabilityRepoClone
	} else if len(parts) >= 5 && parts[0] == "api" && parts[1] == "v3" && parts[2] == "repos" {
		host = lease.Scope.ProviderHost
		repo = normalizeRepo(parts[3] + "/" + parts[4])
		upstreamPath = "/" + strings.Join(parts, "/")
		capability = ghCapability(r.Method, parts[5:])
	} else {
		return "", "", "", "", ErrUnknownMapping
	}
	if host != lease.Scope.ProviderHost || !repoAllowed(lease.Scope.Repositories, repo) {
		return "", "", "", "", ErrForbidden
	}
	if !capabilityAllowed(lease.Capabilities, capability) {
		return "", "", "", "", ErrForbidden
	}
	return host, repo, capability, upstreamPath, nil
}

func RuntimeEnvForLease(env map[string]string) (map[string]string, error) {
	out := map[string]string{}
	for k, v := range env {
		out[k] = v
	}
	token := strings.TrimSpace(out[EnvLeaseToken])
	if token == "" {
		return out, nil
	}
	defer func() {
		delete(out, EnvLeaseToken)
		delete(out, "AGS_TOKEN")
		delete(out, "GITHUB_TOKEN")
		delete(out, "GITEA_TOKEN")
	}()
	base, err := normalizeLocalProxyBase(out[EnvProxyBase])
	if err != nil {
		return out, err
	}
	host := normalizeHost(out[EnvProviderHost])
	if host == "" {
		host = "github.com"
	}
	gitURLPrefix := base + "/_agentteam_git/" + host + "/"
	out["GH_HOST"] = "localhost"
	out["GH_TOKEN"] = token
	out["GH_ENTERPRISE_TOKEN"] = token
	out["GIT_CONFIG_COUNT"] = "2"
	out["GIT_CONFIG_KEY_0"] = "url." + gitURLPrefix + ".insteadOf"
	out["GIT_CONFIG_VALUE_0"] = "https://" + host + "/"
	out["GIT_CONFIG_KEY_1"] = "http." + gitURLPrefix + ".extraHeader"
	out["GIT_CONFIG_VALUE_1"] = "Authorization: Bearer " + token
	return out, nil
}

func RewriteCloneURL(raw, proxyBase string) (string, error) {
	u, err := url.Parse(raw)
	if err != nil {
		return "", err
	}
	host := normalizeHost(u.Host)
	parts := splitPath(u.Path)
	if host == "" || len(parts) < 2 {
		return "", ErrUnknownMapping
	}
	base, err := url.Parse(strings.TrimRight(proxyBase, "/"))
	if err != nil {
		return "", err
	}
	base.Path = joinURLPath(base.Path, "/_agentteam_git/"+host+"/"+parts[0]+"/"+parts[1])
	return base.String(), nil
}
func normalizeLocalProxyBase(raw string) (string, error) {
	raw = strings.TrimRight(strings.TrimSpace(raw), "/")
	if raw == "" {
		raw = "http://localhost"
	}
	u, err := url.Parse(raw)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return "", errors.New("local repository proxy base URL is invalid")
	}
	if u.Scheme != "http" {
		return "", errors.New("local repository proxy base URL must use http")
	}
	if h := strings.ToLower(u.Hostname()); h != "localhost" && h != "127.0.0.1" && h != "::1" {
		return "", errors.New("local repository proxy base URL must be localhost-only")
	}
	u.RawQuery = ""
	u.Fragment = ""
	return strings.TrimRight(u.String(), "/"), nil
}

func sanitizedUpstreamRawQuery(values url.Values) string {
	if len(values) == 0 {
		return ""
	}
	out := url.Values{}
	for k, vs := range values {
		if isSensitiveQueryKey(k) {
			continue
		}
		for _, v := range vs {
			out.Add(k, v)
		}
	}
	return out.Encode()
}

func isSensitiveQueryKey(k string) bool {
	switch strings.ToLower(strings.TrimSpace(k)) {
	case "token", "access_token", "auth", "authorization", "gh_token", "github_token", "agentteam_repository_access_local_proxy_token":
		return true
	default:
		return false
	}
}

func Redact(s string) string { return secretLike.ReplaceAllString(s, "[REDACTED]") }

func ghCapability(method string, rest []string) string {
	if len(rest) > 0 && rest[0] == "pulls" {
		if method == http.MethodGet {
			return CapabilityPullRequestRead
		}
		return CapabilityPullRequestWrite
	}
	if method == http.MethodGet {
		return CapabilityRepoRead
	}
	return CapabilityRepoWrite
}
func randomTokenBytes() ([]byte, error) { b := make([]byte, 32); _, err := rand.Read(b); return b, err }
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
func normalizeHost(s string) string {
	s = strings.ToLower(strings.TrimSpace(s))
	s = strings.TrimPrefix(s, "https://")
	s = strings.TrimPrefix(s, "http://")
	return strings.Trim(s, "/")
}
func normalizeRepo(s string) string {
	s = strings.ToLower(strings.Trim(strings.TrimSpace(s), "/"))
	s = strings.TrimSuffix(s, ".git")
	return s
}
func normalizeRepos(in []string) []string {
	set := map[string]struct{}{}
	for _, r := range in {
		if n := normalizeRepo(r); n != "" {
			set[n] = struct{}{}
		}
	}
	out := make([]string, 0, len(set))
	for r := range set {
		out = append(out, r)
	}
	sort.Strings(out)
	return out
}
func normalizeCapabilities(in []string) []string {
	set := map[string]struct{}{}
	for _, c := range in {
		if c = strings.TrimSpace(c); c != "" {
			set[c] = struct{}{}
		}
	}
	out := make([]string, 0, len(set))
	for c := range set {
		out = append(out, c)
	}
	sort.Strings(out)
	return out
}
func repoAllowed(rs []string, repo string) bool {
	repo = normalizeRepo(repo)
	for _, r := range rs {
		if normalizeRepo(r) == repo {
			return true
		}
	}
	return false
}
func capabilityAllowed(caps []string, want string) bool {
	for _, c := range caps {
		c = strings.TrimSpace(c)
		if c == want {
			return true
		}
		if want == CapabilityRepoClone && c == CapabilityRepoRead {
			return true
		}
	}
	return false
}
func splitPath(p string) []string {
	raw := strings.Split(strings.Trim(path.Clean("/"+p), "/"), "/")
	out := []string{}
	for _, x := range raw {
		if x != "" && x != "." {
			out = append(out, x)
		}
	}
	return out
}
func extractLeaseToken(r *http.Request) string {
	auth := strings.TrimSpace(r.Header.Get("Authorization"))
	if strings.HasPrefix(strings.ToLower(auth), "bearer ") {
		return strings.TrimSpace(auth[7:])
	}
	if u, p, ok := r.BasicAuth(); ok {
		if p != "" {
			return p
		}
		return u
	}
	return ""
}
func copyHeaders(dst, src http.Header) {
	for k, vs := range src {
		if strings.EqualFold(k, "Authorization") {
			continue
		}
		for _, v := range vs {
			dst.Add(k, v)
		}
	}
}
func statusFor(err error) int {
	if errors.Is(err, ErrUnknownMapping) {
		return http.StatusBadRequest
	}
	if errors.Is(err, ErrForbidden) {
		return http.StatusForbidden
	}
	return http.StatusUnauthorized
}
func joinURLPath(a, b string) string {
	return strings.TrimRight(a, "/") + "/" + strings.TrimLeft(b, "/")
}