package ui

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/slanos/turnscale/internal/audit"
	"github.com/slanos/turnscale/internal/config"
	"github.com/slanos/turnscale/internal/identity"
	"github.com/slanos/turnscale/internal/policy"
)

// UI serves the web dashboard for the MCP gateway.
type UI struct {
	cfg      *config.Config
	ident    identity.Identifier
	policy   *policy.Engine
	audit    *audit.Logger
	adminIDs map[string]bool

	mu sync.RWMutex // protects cfg and policy during updates
}

// New creates a new UI handler.
func New(cfg *config.Config, ident identity.Identifier, pol *policy.Engine, aud *audit.Logger, adminIDs map[string]bool) *UI {
	return &UI{
		cfg:      cfg,
		ident:    ident,
		policy:   pol,
		audit:    aud,
		adminIDs: adminIDs,
	}
}

// SetPolicy updates the policy engine (called after config changes).
func (u *UI) SetPolicy(pol *policy.Engine) {
	u.mu.Lock()
	defer u.mu.Unlock()
	u.policy = pol
}

// Servers returns a snapshot copy of the server map (safe for concurrent use).
func (u *UI) Servers() map[string]config.Server {
	u.mu.RLock()
	defer u.mu.RUnlock()
	cp := make(map[string]config.Server, len(u.cfg.Servers))
	for k, v := range u.cfg.Servers {
		cp[k] = v
	}
	return cp
}

// Hostname returns the configured hostname.
func (u *UI) Hostname() string {
	u.mu.RLock()
	defer u.mu.RUnlock()
	return u.cfg.Hostname
}

// ToolInfo holds a discovered tool from a backend server.
type ToolInfo struct {
	Name        string
	Description string
	Denied      bool
}

// ServerStatus holds health, access, and tool info for a server.
type ServerStatus struct {
	Name        string
	URL         string
	Transport   string
	Healthy     bool
	Allowed     bool
	DeniedTools []string
	Tools       []ToolInfo
	ToolCount   int
	Error       string
	Stats       *audit.ServerStat
}

// PolicyView holds display info for a policy.
type PolicyView struct {
	Name      string
	Identity  []string
	Tags      []string
	Allow     []string
	Deny      []string
	DenyTools []string
	IsActive  bool
}

// CallerStat holds a caller's request count for a chart tooltip.
type CallerStat struct {
	Caller string
	Server string
	Count  int
}

// ChartBar represents one bar in a chart.
type ChartBar struct {
	Label   string
	Total   int
	Errors  int
	Denied  int
	Height  int // percentage 0-100
	ErrH    int // error height percentage
	DenyH   int // denied height percentage
	Callers []CallerStat
}

type dashboardData struct {
	Version      string
	Hostname     string
	Caller       *identity.Caller
	IsAdmin      bool
	Servers      []ServerStatus
	Policies     []PolicyView
	RecentAudit  []audit.Row
	TotalTools   int
	HealthyCount int
	AccessCount  int
	Chart         []ChartBar
	ChartMax      int
	ChartTotal    int
	ChartErrors   int
	ChartDenied   int
	HasGrants     bool
	RecIDs        map[int64]bool
	TotalRequests int
}

// HandleDashboard renders the main dashboard page.
func (u *UI) HandleDashboard(w http.ResponseWriter, r *http.Request) {
	caller, err := u.ident.Identify(r)
	if err != nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}

	u.mu.RLock()
	cfg := u.cfg
	pol := u.policy
	u.mu.RUnlock()

	isAdmin := u.adminIDs[caller.UserLogin]
	callerGrant := policy.ParseGrants(caller, u.cfg.Tailnet)
	if callerGrant != nil && callerGrant.Admin {
		isAdmin = true
	}
	hasGrants := callerGrant != nil
	probes := u.probeServers(r.Context(), caller)

	// Build server stats map
	var serverStatsMap map[string]*audit.ServerStat
	if isAdmin {
		if stats, err := u.audit.ServerStats(24); err == nil {
			serverStatsMap = make(map[string]*audit.ServerStat, len(stats))
			for i := range stats {
				serverStatsMap[stats[i].Server] = &stats[i]
			}
		}
	}

	var servers []ServerStatus
	totalTools := 0
	for name, srv := range cfg.Servers {
		probe := probes[name]
		allowed := pol.EvalServer(caller, name) == policy.Allow
		var deniedTools []string
		if allowed {
			deniedTools = pol.DeniedTools(caller, name)
		}

		var tools []ToolInfo
		if allowed && probe.tools != nil {
			deniedPatterns := deniedTools
			for _, t := range probe.tools {
				denied := false
				for _, pattern := range deniedPatterns {
					if matchGlob(pattern, t.Name) {
						denied = true
						break
					}
				}
				tools = append(tools, ToolInfo{
					Name:        t.Name,
					Description: t.Description,
					Denied:      denied,
				})
			}
		}

		toolCount := len(tools)
		totalTools += toolCount

		servers = append(servers, ServerStatus{
			Name:        name,
			URL:         srv.URL,
			Transport:   srv.Transport,
			Healthy:     probe.healthy,
			Allowed:     allowed,
			DeniedTools: deniedTools,
			Tools:       tools,
			ToolCount:   toolCount,
			Error:       probe.err,
			Stats:       serverStatsMap[name],
		})
	}
	sort.Slice(servers, func(i, j int) bool { return servers[i].Name < servers[j].Name })

	healthyCount := 0
	accessCount := 0
	for _, s := range servers {
		if s.Healthy {
			healthyCount++
		}
		if s.Allowed {
			accessCount++
		}
	}

	firstMatch := pol.FirstMatch(caller)
	var policies []PolicyView
	for i, p := range cfg.Policies {
		policies = append(policies, PolicyView{
			Name:      p.Name,
			Identity:  p.Match.Identity,
			Tags:      p.Match.Tags,
			Allow:     p.Allow,
			Deny:      p.Deny,
			DenyTools: p.DenyTools,
			IsActive:  i == firstMatch,
		})
	}

	var recentAudit []audit.Row
	var recIDs map[int64]bool
	var cr *chartResult
	if isAdmin {
		recentAudit, _ = u.audit.Query(audit.QueryParams{Limit: 10})
		cr = u.buildChart(24)
		// Check which audit entries have recordings
		if len(recentAudit) > 0 {
			ids := make([]int64, len(recentAudit))
			for i, r := range recentAudit {
				ids[i] = r.ID
			}
			recIDs = u.audit.RecordingIDs(ids)
		}
	}

	data := dashboardData{
		Version:      "0.1.0",
		Hostname:     cfg.Hostname,
		Caller:       caller,
		IsAdmin:      isAdmin,
		Servers:      servers,
		Policies:     policies,
		RecentAudit:  recentAudit,
		TotalTools:   totalTools,
		HealthyCount: healthyCount,
		AccessCount:  accessCount,
		HasGrants:     hasGrants,
		RecIDs:        recIDs,
		TotalRequests: u.audit.TotalRequests(),
	}
	if cr != nil {
		data.Chart = cr.Bars
		data.ChartMax = cr.MaxVal
		data.ChartTotal = cr.Total
		data.ChartErrors = cr.Errors
		data.ChartDenied = cr.Denied
	}

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	dashboardTmpl.Execute(w, data)
}

// chartResult holds chart bars and aggregate stats.
type chartResult struct {
	Bars    []ChartBar
	MaxVal  int
	Total   int
	Errors  int
	Denied  int
}

// buildChart creates a 24-hour bar chart from audit data.
func (u *UI) buildChart(hours int) *chartResult {
	counts, err := u.audit.HourlyCounts(hours)
	if err != nil || len(counts) == 0 {
		return nil
	}

	// Get per-caller breakdown
	details, _ := u.audit.HourlyBreakdown(hours)
	detailMap := make(map[string][]CallerStat)
	for _, d := range details {
		detailMap[d.Hour] = append(detailMap[d.Hour], CallerStat{
			Caller: d.Caller,
			Server: d.Server,
			Count:  d.Count,
		})
	}

	// Find max for scaling and compute totals
	maxVal := 1
	totalReqs, totalErrs, totalDenied := 0, 0, 0
	for _, c := range counts {
		if c.Total > maxVal {
			maxVal = c.Total
		}
		totalReqs += c.Total
		totalErrs += c.Errors
		totalDenied += c.Denied
	}

	// Build bars indexed by hour
	hourMap := make(map[string]audit.HourlyCount, len(counts))
	for _, c := range counts {
		hourMap[c.Hour] = c
	}

	now := time.Now().UTC()
	bars := make([]ChartBar, hours)
	for i := 0; i < hours; i++ {
		h := now.Add(-time.Duration(hours-1-i) * time.Hour).Truncate(time.Hour)
		key := h.Format("2006-01-02T15:00:00Z")
		c := hourMap[key]
		height := 0
		errH := 0
		denyH := 0
		if c.Total > 0 {
			height = c.Total * 100 / maxVal
			if height < 15 {
				height = 15
			}
			errH = c.Errors * 100 / maxVal
			if c.Errors > 0 && errH < 5 {
				errH = 5
			}
			denyH = c.Denied * 100 / maxVal
			if c.Denied > 0 && denyH < 5 {
				denyH = 5
			}
		}

		// Limit to top 3 caller/server combos to keep tooltips compact
		callers := detailMap[key]
		if len(callers) > 3 {
			callers = callers[:3]
		}
		bars[i] = ChartBar{
			Label:   h.Format("15"),
			Total:   c.Total,
			Errors:  c.Errors,
			Denied:  c.Denied,
			Height:  height,
			ErrH:    errH,
			DenyH:   denyH,
			Callers: callers,
		}
	}
	return &chartResult{Bars: bars, MaxVal: maxVal, Total: totalReqs, Errors: totalErrs, Denied: totalDenied}
}

var validServerName = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]*$`)

// HandleAddServer handles POST /ui/servers to add a new server.
func (u *UI) HandleAddServer(w http.ResponseWriter, r *http.Request) {
	caller, err := u.ident.Identify(r)
	if err != nil || !u.adminIDs[caller.UserLogin] {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}

	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", http.StatusBadRequest)
		return
	}

	name := strings.TrimSpace(r.FormValue("name"))
	url := strings.TrimSpace(r.FormValue("url"))
	transport := strings.TrimSpace(r.FormValue("transport"))
	if name == "" || url == "" {
		http.Error(w, "name and url required", http.StatusBadRequest)
		return
	}
	if !validServerName.MatchString(name) {
		http.Error(w, "invalid server name: use lowercase letters, numbers, hyphens, underscores", http.StatusBadRequest)
		return
	}
	if transport == "" {
		transport = "streamable-http"
	}

	u.mu.Lock()
	u.cfg.Servers[name] = config.Server{URL: url, Transport: transport}
	if err := u.cfg.Save(); err != nil {
		slog.Error("config save failed", "error", err)
	}
	u.mu.Unlock()

	http.Redirect(w, r, "/ui/", http.StatusSeeOther)
}

// HandleDeleteServer handles POST /ui/servers/delete to remove a server.
func (u *UI) HandleDeleteServer(w http.ResponseWriter, r *http.Request) {
	caller, err := u.ident.Identify(r)
	if err != nil || !u.adminIDs[caller.UserLogin] {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}

	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", http.StatusBadRequest)
		return
	}

	name := strings.TrimSpace(r.FormValue("name"))
	if name == "" {
		http.Error(w, "name required", http.StatusBadRequest)
		return
	}

	u.mu.Lock()
	delete(u.cfg.Servers, name)
	if err := u.cfg.Save(); err != nil {
		slog.Error("config save failed", "error", err)
	}
	u.mu.Unlock()

	http.Redirect(w, r, "/ui/", http.StatusSeeOther)
}

// HandleEditServer handles POST /ui/servers/edit to update a server.
func (u *UI) HandleEditServer(w http.ResponseWriter, r *http.Request) {
	caller, err := u.ident.Identify(r)
	if err != nil || !u.adminIDs[caller.UserLogin] {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}

	if err := r.ParseForm(); err != nil {
		http.Error(w, "bad form", http.StatusBadRequest)
		return
	}

	name := strings.TrimSpace(r.FormValue("name"))
	url := strings.TrimSpace(r.FormValue("url"))
	transport := strings.TrimSpace(r.FormValue("transport"))
	if name == "" || url == "" {
		http.Error(w, "name and url required", http.StatusBadRequest)
		return
	}
	if transport == "" {
		transport = "streamable-http"
	}

	u.mu.Lock()
	u.cfg.Servers[name] = config.Server{URL: url, Transport: transport}
	if err := u.cfg.Save(); err != nil {
		slog.Error("config save failed", "error", err)
	}
	u.mu.Unlock()

	http.Redirect(w, r, "/ui/", http.StatusSeeOther)
}

// HandleSession shows a session recording detail page.
func (u *UI) HandleSession(w http.ResponseWriter, r *http.Request) {
	caller, err := u.ident.Identify(r)
	if err != nil {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}
	isAdmin := u.adminIDs[caller.UserLogin]
	if !isAdmin {
		if g := policy.ParseGrants(caller, u.cfg.Tailnet); g != nil {
			isAdmin = g.Admin
		}
	}
	if !isAdmin {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}

	idStr := r.PathValue("id")
	id, err := strconv.ParseInt(idStr, 10, 64)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	rec, err := u.audit.GetRecording(id)
	if err != nil {
		http.NotFound(w, r)
		return
	}

	// Pretty-print JSON bodies
	rec.Request = prettyJSON(rec.Request)
	rec.Response = prettyJSON(rec.Response)

	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	sessionTmpl.Execute(w, rec)
}

func prettyJSON(s string) string {
	var buf bytes.Buffer
	if err := json.Indent(&buf, []byte(s), "", "  "); err != nil {
		return s
	}
	return buf.String()
}


// probeResult holds the result of probing a single backend server.
type probeResult struct {
	healthy bool
	tools   []mcpTool
	err     string
}

// mcpTool is a tool discovered from a backend's tools/list response.
type mcpTool struct {
	Name        string `json:"name"`
	Description string `json:"description"`
}

// probeServers probes all backend servers concurrently for health and tools.
func (u *UI) probeServers(ctx context.Context, caller *identity.Caller) map[string]probeResult {
	u.mu.RLock()
	servers := u.cfg.Servers
	u.mu.RUnlock()

	results := make(map[string]probeResult, len(servers))
	var mu sync.Mutex
	var wg sync.WaitGroup

	for name, srv := range servers {
		wg.Add(1)
		go func(name, url string) {
			defer wg.Done()
			result := probeBackend(ctx, url)
			mu.Lock()
			results[name] = result
			mu.Unlock()
		}(name, srv.URL)
	}
	wg.Wait()
	return results
}

// probeBackend sends MCP initialize + tools/list to a backend and returns health + tools.
func probeBackend(ctx context.Context, url string) probeResult {
	client := &http.Client{Timeout: 5 * time.Second}

	initReq := jsonRPCRequest{
		JSONRPC: "2.0",
		ID:      1,
		Method:  "initialize",
		Params: map[string]any{
			"protocolVersion": "2025-03-26",
			"capabilities":   map[string]any{},
			"clientInfo": map[string]string{
				"name":    "turnscale-ui",
				"version": "0.1.0",
			},
		},
	}

	initBody, _ := json.Marshal(initReq)
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(initBody))
	if err != nil {
		return probeResult{err: err.Error()}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")

	resp, err := client.Do(req)
	if err != nil {
		return probeResult{err: err.Error()}
	}
	defer resp.Body.Close()
	io.ReadAll(resp.Body)

	if resp.StatusCode >= 400 {
		return probeResult{healthy: false, err: fmt.Sprintf("initialize: %s", resp.Status)}
	}

	sessionID := resp.Header.Get("Mcp-Session-Id")

	notifReq := jsonRPCRequest{JSONRPC: "2.0", Method: "notifications/initialized"}
	notifBody, _ := json.Marshal(notifReq)
	notifHTTP, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(notifBody))
	if err == nil {
		notifHTTP.Header.Set("Content-Type", "application/json")
		notifHTTP.Header.Set("Accept", "application/json, text/event-stream")
		if sessionID != "" {
			notifHTTP.Header.Set("Mcp-Session-Id", sessionID)
		}
		if r, err := client.Do(notifHTTP); err == nil {
			io.ReadAll(r.Body)
			r.Body.Close()
		}
	}

	toolsReq := jsonRPCRequest{JSONRPC: "2.0", ID: 2, Method: "tools/list"}
	toolsBody, _ := json.Marshal(toolsReq)
	toolsHTTP, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(toolsBody))
	if err != nil {
		return probeResult{healthy: true, err: "tools/list request failed"}
	}
	toolsHTTP.Header.Set("Content-Type", "application/json")
	toolsHTTP.Header.Set("Accept", "application/json, text/event-stream")
	if sessionID != "" {
		toolsHTTP.Header.Set("Mcp-Session-Id", sessionID)
	}

	toolsResp, err := client.Do(toolsHTTP)
	if err != nil {
		return probeResult{healthy: true, err: "tools/list: " + err.Error()}
	}
	defer toolsResp.Body.Close()

	toolsRespBody, err := io.ReadAll(toolsResp.Body)
	if err != nil {
		return probeResult{healthy: true, err: "reading tools/list response"}
	}

	tools := parseToolsList(toolsRespBody)
	sort.Slice(tools, func(i, j int) bool { return tools[i].Name < tools[j].Name })

	return probeResult{healthy: true, tools: tools}
}

type jsonRPCRequest struct {
	JSONRPC string `json:"jsonrpc"`
	ID      any    `json:"id,omitempty"`
	Method  string `json:"method"`
	Params  any    `json:"params,omitempty"`
}

func parseToolsList(body []byte) []mcpTool {
	var resp struct {
		Result struct {
			Tools []mcpTool `json:"tools"`
		} `json:"result"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil
	}
	return resp.Result.Tools
}

func matchGlob(pattern, name string) bool {
	if idx := len(pattern) - 1; idx >= 0 && pattern[idx] == '*' {
		return strings.HasPrefix(name, pattern[:idx])
	}
	return pattern == name
}

func initial(s string) string {
	if len(s) > 0 {
		return strings.ToUpper(string(s[0]))
	}
	return "?"
}