package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/ongridio/ongrid/internal/pkg/tracequery"
)

// ToolNameQueryTraceQL is the stable wire name the LLM sees for the TraceQL
// tool.
const ToolNameQueryTraceQL = "query_traceql"

// QueryTraceQLDescription is the single-sentence description shown to the
// LLM. Phrased to point the model at trace search whenever metrics + logs
// don't pin down which request is slow.
const QueryTraceQLDescription = "Run a TraceQL search against Tempo. " +
	"Use this to find traces by service / operation / latency / status. " +
	"Returns trace summaries (id, service, root span name, duration, span count)."

// QueryTraceQLSchema is the JSON Schema of the tool's argument object.
// query may be empty when service/operation/duration filters are supplied —
// Tempo's tag-mode search handles that case.
var QueryTraceQLSchema = json.RawMessage(`{
  "type": "object",
  "properties": {
    "query": {
      "type": "string",
      "description": "TraceQL expression. May be empty if service/operation/duration filters are given."
    },
    "service": {
      "type": "string",
      "description": "Filter by resource.service.name (tag mode)."
    },
    "operation": {
      "type": "string",
      "description": "Filter by span name (tag mode)."
    },
    "start": {
      "type": "string",
      "description": "RFC3339 start time. Defaults to now-1h."
    },
    "end": {
      "type": "string",
      "description": "RFC3339 end time. Defaults to now."
    },
    "limit": {
      "type": "integer",
      "minimum": 1,
      "maximum": 1000,
      "description": "Max trace summaries to return (default 50)."
    },
    "min_duration": {
      "type": "string",
      "description": "Minimum trace duration as a Go duration string (e.g. \"100ms\", \"2s\")."
    },
    "max_duration": {
      "type": "string",
      "description": "Maximum trace duration as a Go duration string."
    }
  }
}`)

// QueryTraceQLArgs is the typed form of QueryTraceQLSchema.
type QueryTraceQLArgs struct {
	Query       string `json:"query,omitempty"`
	Service     string `json:"service,omitempty"`
	Operation   string `json:"operation,omitempty"`
	Start       string `json:"start,omitempty"`
	End         string `json:"end,omitempty"`
	Limit       int    `json:"limit,omitempty"`
	MinDuration string `json:"min_duration,omitempty"`
	MaxDuration string `json:"max_duration,omitempty"`
}

// queryTraceqlCallTimeout caps how long a single dispatch may wait. Mirrors
// query_promql for symmetry across signal types.
const queryTraceqlCallTimeout = 30 * time.Second

// executeQueryTraceQL runs the TraceQL search and hands the raw Tempo
// response back to the LLM via ResultJSON.
func (r *Registry) executeQueryTraceQL(ctx context.Context, args json.RawMessage) (ExecuteResult, error) {
	if r.traceQuery == nil {
		// Should not happen — when traceQuery is nil at NewRegistry the
		// tool is never registered. Defensive guard.
		return ExecuteResult{}, fmt.Errorf("query_traceql: trace query client not configured")
	}
	var in QueryTraceQLArgs
	if err := json.Unmarshal(args, &in); err != nil {
		return ExecuteResult{}, fmt.Errorf("query_traceql: bad args: %w", err)
	}

	// Require *some* filter — either a TraceQL query, a tag (service /
	// operation), or a duration bound. An unfiltered Tempo search is too
	// expensive to dump on the LLM by accident.
	if strings.TrimSpace(in.Query) == "" &&
		strings.TrimSpace(in.Service) == "" &&
		strings.TrimSpace(in.Operation) == "" &&
		strings.TrimSpace(in.MinDuration) == "" &&
		strings.TrimSpace(in.MaxDuration) == "" {
		return ExecuteResult{}, fmt.Errorf("query_traceql: at least one of query/service/operation/min_duration/max_duration required")
	}

	end := time.Now()
	start := end.Add(-time.Hour)
	if in.End != "" {
		t, err := time.Parse(time.RFC3339, in.End)
		if err != nil {
			return ExecuteResult{}, fmt.Errorf("query_traceql: parse end: %w", err)
		}
		end = t
	}
	if in.Start != "" {
		t, err := time.Parse(time.RFC3339, in.Start)
		if err != nil {
			return ExecuteResult{}, fmt.Errorf("query_traceql: parse start: %w", err)
		}
		start = t
	} else if in.End != "" {
		start = end.Add(-time.Hour)
	}

	limit := in.Limit
	if limit <= 0 {
		limit = 50
	}

	var minDur, maxDur time.Duration
	if in.MinDuration != "" {
		d, err := time.ParseDuration(in.MinDuration)
		if err != nil {
			return ExecuteResult{}, fmt.Errorf("query_traceql: parse min_duration: %w", err)
		}
		minDur = d
	}
	if in.MaxDuration != "" {
		d, err := time.ParseDuration(in.MaxDuration)
		if err != nil {
			return ExecuteResult{}, fmt.Errorf("query_traceql: parse max_duration: %w", err)
		}
		maxDur = d
	}

	tags := map[string]string{}
	if in.Service != "" {
		tags["service.name"] = in.Service
	}
	if in.Operation != "" {
		tags["name"] = in.Operation
	}
	// Tempo's SearchTraces ignores Tags when Query is set — that's the
	// desired precedence. We forward both and let the client choose.
	if len(tags) == 0 {
		tags = nil
	}

	callCtx, cancel := context.WithTimeout(ctx, queryTraceqlCallTimeout)
	defer cancel()

	res, err := r.traceQuery.SearchTraces(callCtx, tracequery.SearchOptions{
		Query:       in.Query,
		Tags:        tags,
		Limit:       limit,
		Start:       start,
		End:         end,
		MinDuration: minDur,
		MaxDuration: maxDur,
	})
	if err != nil {
		return ExecuteResult{}, fmt.Errorf("query_traceql: dispatch: %w", err)
	}
	out, err := json.Marshal(res)
	if err != nil {
		return ExecuteResult{}, fmt.Errorf("query_traceql: marshal response: %w", err)
	}
	return ExecuteResult{ResultJSON: out}, nil
}

// TraceQuerier is the narrow surface the query_traceql executor needs from
// the tracequery client. Declared here so tests can inject a fake.
//
// NOTE: this interface is what r.traceQuery is typed as. The concrete
// *tracequery.Client satisfies it.
type TraceQuerier interface {
	SearchTraces(ctx context.Context, opts tracequery.SearchOptions) (*tracequery.SearchResult, error)
}
