// Package mca provides OpenTelemetry instrumentation for ML models.
//
// The MCA (Model Collector Agent) SDK enables healthcare ML models to emit
// metrics, logs, and distributed traces through OpenTelemetry.
package mca

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"sync"
	"time"

	"github.com/baptisthealth/mca-sdk-go/config"
	"github.com/baptisthealth/mca-sdk-go/internal/telemetry"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/log"
	"go.opentelemetry.io/otel/metric"
	sdklog "go.opentelemetry.io/otel/sdk/log"
	sdkmetric "go.opentelemetry.io/otel/sdk/metric"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/trace"
)

const (
	// MaxCustomAttributes is the maximum number of custom attributes per prediction
	MaxCustomAttributes = 50
	// MaxAttributeValueLength is the maximum length for attribute string values
	MaxAttributeValueLength = 256
)

// ErrClientClosed is returned when operations are attempted on a closed client
var ErrClientClosed = errors.New("client has been shut down")

// Client is the main entry point for MCA SDK instrumentation.
// It manages OpenTelemetry providers and provides methods for recording
// model predictions with metrics, logs, and traces.
//
// Client is goroutine-safe and can be used concurrently from multiple goroutines.
type Client struct {
	config         *config.Config
	resource       *resource.Resource
	meterProvider  *sdkmetric.MeterProvider
	tracerProvider *sdktrace.TracerProvider
	loggerProvider *sdklog.LoggerProvider
	meter          metric.Meter
	tracer         trace.Tracer
	logger         log.Logger

	// Metrics instruments
	predictionCounter metric.Int64Counter
	latencyHistogram  metric.Float64Histogram

	// Shutdown management
	shutdownOnce sync.Once
	shutdownErr  error
	isShutdown   bool
	mu           sync.RWMutex
}

// NewClient creates and initializes a new MCA Client with a context.
//
// The context is used for provider initialization and has a timeout to prevent
// indefinite hangs if the collector is unreachable.
//
// Example:
//
//	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
//	defer cancel()
//	cfg := &config.Config{
//	    ServiceName: "readmission-model",
//	    ModelID:     "mdl-001",
//	    TeamName:    "clinical-ai",
//	}
//	client, err := mca.NewClient(ctx, cfg)
//	if err != nil {
//	    return err
//	}
//	defer client.Shutdown()
func NewClient(ctx context.Context, cfg *config.Config) (*Client, error) {
	// Apply defaults and validate configuration
	cfg.ApplyDefaults()
	if err := cfg.Validate(); err != nil {
		return nil, err
	}

	// Create resource with model attributes
	res, err := telemetry.CreateResource(
		cfg.ServiceName,
		cfg.ModelID,
		cfg.ModelVersion,
		cfg.TeamName,
		cfg.ModelType,
		cfg.ModelCategory,
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create resource: %w", err)
	}

	// Parse endpoint properly using net/url
	parsedURL, err := url.Parse(cfg.CollectorEndpoint)
	if err != nil {
		return nil, fmt.Errorf("invalid collector endpoint: %w", err)
	}
	endpoint := parsedURL.Host
	if endpoint == "" {
		endpoint = parsedURL.Path
	}

	// Create timeout context for initialization if not already bounded
	initCtx := ctx
	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
		var cancel context.CancelFunc
		initCtx, cancel = context.WithTimeout(ctx, 30*time.Second)
		defer cancel()
	}

	// Setup OpenTelemetry providers
	meterProvider, err := telemetry.SetupMeterProvider(initCtx, endpoint, res, cfg.TLSInsecure)
	if err != nil {
		return nil, fmt.Errorf("failed to setup meter provider: %w", err)
	}

	tracerProvider, err := telemetry.SetupTracerProvider(initCtx, endpoint, res, cfg.TLSInsecure)
	if err != nil {
		// Rollback with timeout context
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = meterProvider.Shutdown(rollbackCtx)
		cancel()
		return nil, fmt.Errorf("failed to setup tracer provider: %w", err)
	}

	loggerProvider, err := telemetry.SetupLoggerProvider(initCtx, endpoint, res, cfg.TLSInsecure)
	if err != nil {
		// Rollback with timeout context
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = meterProvider.Shutdown(rollbackCtx)
		_ = tracerProvider.Shutdown(rollbackCtx)
		cancel()
		return nil, fmt.Errorf("failed to setup logger provider: %w", err)
	}

	// Create meter, tracer, and logger from providers
	meter := meterProvider.Meter("mca-sdk-go")
	tracer := tracerProvider.Tracer("mca-sdk-go")
	logger := loggerProvider.Logger("mca-sdk-go")

	// Create metric instruments
	predictionCounter, err := meter.Int64Counter(
		"model.predictions_total",
		metric.WithDescription("Total number of model predictions"),
		metric.WithUnit("{prediction}"),
	)
	if err != nil {
		// Rollback with timeout context
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = meterProvider.Shutdown(rollbackCtx)
		_ = tracerProvider.Shutdown(rollbackCtx)
		_ = loggerProvider.Shutdown(rollbackCtx)
		cancel()
		return nil, fmt.Errorf("failed to create prediction counter: %w", err)
	}

	latencyHistogram, err := meter.Float64Histogram(
		"model.latency_seconds",
		metric.WithDescription("Model prediction latency"),
		metric.WithUnit("s"),
	)
	if err != nil {
		// Rollback with timeout context
		rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		_ = meterProvider.Shutdown(rollbackCtx)
		_ = tracerProvider.Shutdown(rollbackCtx)
		_ = loggerProvider.Shutdown(rollbackCtx)
		cancel()
		return nil, fmt.Errorf("failed to create latency histogram: %w", err)
	}

	client := &Client{
		config:            cfg,
		resource:          res,
		meterProvider:     meterProvider,
		tracerProvider:    tracerProvider,
		loggerProvider:    loggerProvider,
		meter:             meter,
		tracer:            tracer,
		logger:            logger,
		predictionCounter: predictionCounter,
		latencyHistogram:  latencyHistogram,
		isShutdown:        false,
	}

	return client, nil
}

// RecordPrediction records a model prediction event with metrics, logs, and traces.
//
// This method creates a trace span, increments the prediction counter, records latency
// if provided, and emits a structured log with the prediction details.
//
// The context is used for trace propagation and cancellation. If the context
// contains an active span, the prediction will be correlated with that span.
//
// Custom attributes are validated and limited to 50 attributes with 256 char values.
//
// Example:
//
//	latency := 0.15
//	pred := &mca.Prediction{
//	    Latency:      &latency,
//	    PredictionID: "pred-123",
//	    Attributes: map[string]string{
//	        "model_version": "2.0",
//	        "input_size":    "1024",
//	    },
//	}
//	err := client.RecordPrediction(ctx, pred)
func (c *Client) RecordPrediction(ctx context.Context, pred *Prediction) error {
	c.mu.RLock()
	defer c.mu.RUnlock() // Hold read lock for entire method to prevent race with Shutdown

	if c.isShutdown {
		return ErrClientClosed
	}

	// Create trace span for correlation
	ctx, span := c.tracer.Start(ctx, "record_prediction",
		trace.WithSpanKind(trace.SpanKindInternal),
	)
	defer span.End()

	// Build system attributes first (will take precedence)
	attrs := []attribute.KeyValue{
		attribute.String("model.id", c.config.ModelID),
		attribute.String("model.version", c.config.ModelVersion),
		attribute.String("model.type", c.config.ModelType),
	}

	if pred.PredictionID != "" {
		attrs = append(attrs, attribute.String("prediction.id", pred.PredictionID))
		span.SetAttributes(attribute.String("prediction.id", pred.PredictionID))
	}

	// Validate and add custom attributes (before system attrs to allow system override)
	if len(pred.Attributes) > 0 {
		if len(pred.Attributes) > MaxCustomAttributes {
			c.logger.Emit(ctx, log.Record{
				Severity: log.SeverityWarn,
				Body:     log.StringValue(fmt.Sprintf("Prediction has %d attributes, maximum %d allowed - excess ignored", len(pred.Attributes), MaxCustomAttributes)),
			})
		}

		count := 0
		for k, v := range pred.Attributes {
			if count >= MaxCustomAttributes {
				break
			}
			// Truncate long values
			if len(v) > MaxAttributeValueLength {
				v = v[:MaxAttributeValueLength]
			}
			// Add user attributes before system attributes
			attrs = append([]attribute.KeyValue{attribute.String(k, v)}, attrs...)
			count++
		}
	}

	attrSet := attribute.NewSet(attrs...)

	// Increment prediction counter
	c.predictionCounter.Add(ctx, 1, metric.WithAttributeSet(attrSet))

	// Record latency histogram if latency is provided (nil means not provided)
	if pred.Latency != nil {
		c.latencyHistogram.Record(ctx, *pred.Latency, metric.WithAttributeSet(attrSet))
		span.SetAttributes(attribute.Float64("latency_seconds", *pred.Latency))
		span.AddEvent("latency_recorded", trace.WithAttributes(
			attribute.Float64("value", *pred.Latency),
		))
	}

	// Emit structured log with proper type handling
	logRecord := log.Record{}
	logRecord.SetTimestamp(time.Now())
	logRecord.SetBody(log.StringValue("Model prediction recorded"))
	logRecord.SetSeverity(log.SeverityInfo)

	// Add trace context to log for correlation
	spanCtx := span.SpanContext()
	if spanCtx.IsValid() {
		logRecord.AddAttributes(
			log.String("trace_id", spanCtx.TraceID().String()),
			log.String("span_id", spanCtx.SpanID().String()),
		)
	}

	// Add attributes to log record with proper type handling
	for _, attr := range attrs {
		switch attr.Value.Type() {
		case attribute.BOOL:
			logRecord.AddAttributes(log.Bool(string(attr.Key), attr.Value.AsBool()))
		case attribute.INT64:
			logRecord.AddAttributes(log.Int64(string(attr.Key), attr.Value.AsInt64()))
		case attribute.FLOAT64:
			logRecord.AddAttributes(log.Float64(string(attr.Key), attr.Value.AsFloat64()))
		case attribute.STRING:
			logRecord.AddAttributes(log.String(string(attr.Key), attr.Value.AsString()))
		default:
			logRecord.AddAttributes(log.String(string(attr.Key), attr.Value.AsString()))
		}
	}

	if pred.Latency != nil {
		logRecord.AddAttributes(log.Float64("latency_seconds", *pred.Latency))
	}

	c.logger.Emit(ctx, logRecord)

	// Mark span as successful
	span.SetStatus(codes.Ok, "Prediction recorded")

	return nil
}

// Shutdown gracefully shuts down all OpenTelemetry providers.
//
// This method flushes any pending telemetry data and releases resources.
// It is safe to call multiple times (subsequent calls are no-ops).
//
// Shutdown uses a 30-second timeout for graceful shutdown. If shutdown
// takes longer than 30 seconds, it returns an error.
//
// This method should be called with defer immediately after creating the client:
//
//	client, err := mca.NewClient(ctx, cfg)
//	if err != nil {
//	    return err
//	}
//	defer client.Shutdown()
func (c *Client) Shutdown() error {
	c.shutdownOnce.Do(func() {
		c.mu.Lock()
		defer c.mu.Unlock()

		// Mark as shutdown to prevent new operations
		c.isShutdown = true

		// Create context with timeout for shutdown
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()

		var errs []error

		// Shutdown meter provider
		if err := c.meterProvider.Shutdown(ctx); err != nil {
			errs = append(errs, fmt.Errorf("meter provider shutdown: %w", err))
		}

		// Shutdown tracer provider
		if err := c.tracerProvider.Shutdown(ctx); err != nil {
			errs = append(errs, fmt.Errorf("tracer provider shutdown: %w", err))
		}

		// Shutdown logger provider
		if err := c.loggerProvider.Shutdown(ctx); err != nil {
			errs = append(errs, fmt.Errorf("logger provider shutdown: %w", err))
		}

		// Combine errors using errors.Join (Go 1.20+)
		if len(errs) > 0 {
			c.shutdownErr = errors.Join(errs...)
		}
	})

	return c.shutdownErr
}

// Config returns a copy of the client configuration.
func (c *Client) Config() config.Config {
	return *c.config
}
