package db

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"os"
	"time"

	_ "github.com/lib/pq"
)

// DB wraps the database connection
type DB struct {
	conn *sql.DB
}

// New creates a new database connection to memory-db (K8s CNPG)
// Expected DSN format: postgresql://app:password@memory-db-rw.poimen.svc.cluster.local:5432/memory?sslmode=disable
func New(dsn string) (*DB, error) {
	if dsn == "" {
		// Fallback: try to construct from K8s env vars
		host := os.Getenv("DATABASE_HOST")
		port := os.Getenv("DATABASE_PORT")
		name := os.Getenv("DATABASE_NAME")
		user := os.Getenv("DATABASE_USER")
		password := os.Getenv("DATABASE_PASSWORD")
		
		if host != "" && port != "" && name != "" && user != "" && password != "" {
			dsn = fmt.Sprintf("postgresql://%s:%s@%s:%s/%s?sslmode=disable", 
				user, password, host, port, name)
		} else {
			return nil, fmt.Errorf("DATABASE_URL or K8s env vars (DATABASE_HOST, DATABASE_PORT, DATABASE_NAME, DATABASE_USER, DATABASE_PASSWORD) required")
		}
	}

	conn, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open database: %w", err)
	}

	// Test connection
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := conn.PingContext(ctx); err != nil {
		return nil, fmt.Errorf("failed to ping database: %w", err)
	}

	// Set connection pool settings
	conn.SetMaxOpenConns(25)
	conn.SetMaxIdleConns(5)
	conn.SetConnMaxLifetime(5 * time.Minute)

	return &DB{conn: conn}, nil
}

// Close closes the database connection
func (db *DB) Close() error {
	return db.conn.Close()
}

// SaveWorkflow saves or updates a workflow with canvas
func (db *DB) SaveWorkflow(ctx context.Context, wf *Workflow) error {
	query := `
		INSERT INTO workflows (id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
		ON CONFLICT (id) DO UPDATE SET
			name = $3,
			description = $4,
			status = $5,
			version = $6,
			nodes = $7,
			edges = $8,
			updated_at = $11
	`

	_, err := db.conn.ExecContext(ctx, query,
		wf.ID,
		wf.CustomerID,
		wf.Name,
		wf.Description,
		wf.Status,
		wf.Version,
		wf.Nodes,
		wf.Edges,
		wf.CreatedBy,
		wf.CreatedAt,
		wf.UpdatedAt,
	)

	return err
}

// SaveCanvasUpdate saves canvas (nodes + edges) for a workflow
func (db *DB) SaveCanvasUpdate(ctx context.Context, workflowID, customerID string, canvas *Canvas) error {
	nodesJSON, err := json.Marshal(canvas.Nodes)
	if err != nil {
		return fmt.Errorf("failed to marshal nodes: %w", err)
	}

	edgesJSON, err := json.Marshal(canvas.Edges)
	if err != nil {
		return fmt.Errorf("failed to marshal edges: %w", err)
	}

	query := `
		UPDATE workflows
		SET nodes = $1, edges = $2, updated_at = now()
		WHERE id = $3 AND customer_id = $4
	`

	result, err := db.conn.ExecContext(ctx, query, nodesJSON, edgesJSON, workflowID, customerID)
	if err != nil {
		return fmt.Errorf("failed to update canvas: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return err
	}

	if rows == 0 {
		return fmt.Errorf("workflow not found: %s", workflowID)
	}

	return nil
}

// FetchWorkflow retrieves a workflow by ID
func (db *DB) FetchWorkflow(ctx context.Context, workflowID, customerID string) (*Workflow, error) {
	query := `
		SELECT id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at, last_executed_at
		FROM workflows
		WHERE id = $1 AND customer_id = $2
	`

	wf := &Workflow{}
	err := db.conn.QueryRowContext(ctx, query, workflowID, customerID).Scan(
		&wf.ID,
		&wf.CustomerID,
		&wf.Name,
		&wf.Description,
		&wf.Status,
		&wf.Version,
		&wf.Nodes,
		&wf.Edges,
		&wf.CreatedBy,
		&wf.CreatedAt,
		&wf.UpdatedAt,
		&wf.LastExecutedAt,
	)

	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("workflow not found: %s", workflowID)
		}
		return nil, fmt.Errorf("failed to fetch workflow: %w", err)
	}

	return wf, nil
}

// FetchCanvas retrieves canvas (nodes + edges) for a workflow
func (db *DB) FetchCanvas(ctx context.Context, workflowID, customerID string) (*Canvas, error) {
	query := `
		SELECT nodes, edges
		FROM workflows
		WHERE id = $1 AND customer_id = $2
	`

	var nodesJSON, edgesJSON []byte
	err := db.conn.QueryRowContext(ctx, query, workflowID, customerID).Scan(&nodesJSON, &edgesJSON)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("workflow not found: %s", workflowID)
		}
		return nil, fmt.Errorf("failed to fetch canvas: %w", err)
	}

	var nodes []WorkflowNode
	var edges []WorkflowEdge

	if err := json.Unmarshal(nodesJSON, &nodes); err != nil {
		return nil, fmt.Errorf("failed to unmarshal nodes: %w", err)
	}

	if err := json.Unmarshal(edgesJSON, &edges); err != nil {
		return nil, fmt.Errorf("failed to unmarshal edges: %w", err)
	}

	return &Canvas{Nodes: nodes, Edges: edges}, nil
}

// ListWorkflows retrieves all workflows for a customer
func (db *DB) ListWorkflows(ctx context.Context, customerID string, limit, offset int) ([]Workflow, error) {
	query := `
		SELECT id, customer_id, name, description, status, version, nodes, edges, created_by, created_at, updated_at, last_executed_at
		FROM workflows
		WHERE customer_id = $1
		ORDER BY created_at DESC
		LIMIT $2 OFFSET $3
	`

	rows, err := db.conn.QueryContext(ctx, query, customerID, limit, offset)
	if err != nil {
		return nil, fmt.Errorf("failed to list workflows: %w", err)
	}
	defer rows.Close()

	var workflows []Workflow
	for rows.Next() {
		wf := Workflow{}
		err := rows.Scan(
			&wf.ID,
			&wf.CustomerID,
			&wf.Name,
			&wf.Description,
			&wf.Status,
			&wf.Version,
			&wf.Nodes,
			&wf.Edges,
			&wf.CreatedBy,
			&wf.CreatedAt,
			&wf.UpdatedAt,
			&wf.LastExecutedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan workflow: %w", err)
		}
		workflows = append(workflows, wf)
	}

	return workflows, rows.Err()
}

// DeleteWorkflow deletes a workflow
func (db *DB) DeleteWorkflow(ctx context.Context, workflowID, customerID string) error {
	query := `
		DELETE FROM workflows
		WHERE id = $1 AND customer_id = $2
	`

	result, err := db.conn.ExecContext(ctx, query, workflowID, customerID)
	if err != nil {
		return fmt.Errorf("failed to delete workflow: %w", err)
	}

	rows, err := result.RowsAffected()
	if err != nil {
		return err
	}

	if rows == 0 {
		return fmt.Errorf("workflow not found: %s", workflowID)
	}

	return nil
}

// SaveExecution saves a workflow execution record
func (db *DB) SaveExecution(ctx context.Context, exec *WorkflowExecution) error {
	query := `
		INSERT INTO workflow_executions (id, workflow_id, customer_id, temporal_id, status, inputs, outputs, started_at, completed_at, duration_ms, error_message, error_count)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
		ON CONFLICT (id) DO UPDATE SET
			status = $5,
			outputs = $7,
			completed_at = $9,
			duration_ms = $10,
			error_message = $11,
			error_count = $12
	`

	_, err := db.conn.ExecContext(ctx, query,
		exec.ID,
		exec.WorkflowID,
		exec.CustomerID,
		exec.TemporalID,
		exec.Status,
		exec.Inputs,
		exec.Outputs,
		exec.StartedAt,
		exec.CompletedAt,
		exec.DurationMs,
		exec.ErrorMessage,
		exec.ErrorCount,
	)

	return err
}

// FetchExecution retrieves a workflow execution
func (db *DB) FetchExecution(ctx context.Context, executionID string) (*WorkflowExecution, error) {
	query := `
		SELECT id, workflow_id, customer_id, temporal_id, status, inputs, outputs, started_at, completed_at, duration_ms, error_message, error_count
		FROM workflow_executions
		WHERE id = $1
	`

	exec := &WorkflowExecution{}
	err := db.conn.QueryRowContext(ctx, query, executionID).Scan(
		&exec.ID,
		&exec.WorkflowID,
		&exec.CustomerID,
		&exec.TemporalID,
		&exec.Status,
		&exec.Inputs,
		&exec.Outputs,
		&exec.StartedAt,
		&exec.CompletedAt,
		&exec.DurationMs,
		&exec.ErrorMessage,
		&exec.ErrorCount,
	)

	if err != nil {
		if err == sql.ErrNoRows {
			return nil, fmt.Errorf("execution not found: %s", executionID)
		}
		return nil, fmt.Errorf("failed to fetch execution: %w", err)
	}

	return exec, nil
}

// SaveExecutionLog saves an activity log entry
func (db *DB) SaveExecutionLog(ctx context.Context, log *ExecutionLog) error {
	query := `
		INSERT INTO execution_logs (execution_id, node_id, activity_name, level, message, metadata, logged_at)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
	`

	_, err := db.conn.ExecContext(ctx, query,
		log.ExecutionID,
		log.NodeID,
		log.ActivityName,
		log.Level,
		log.Message,
		log.Metadata,
		log.LoggedAt,
	)

	return err
}

// FetchExecutionLogs retrieves all logs for an execution
func (db *DB) FetchExecutionLogs(ctx context.Context, executionID string) ([]ExecutionLog, error) {
	query := `
		SELECT id, execution_id, node_id, activity_name, level, message, metadata, logged_at
		FROM execution_logs
		WHERE execution_id = $1
		ORDER BY logged_at ASC
	`

	rows, err := db.conn.QueryContext(ctx, query, executionID)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch execution logs: %w", err)
	}
	defer rows.Close()

	var logs []ExecutionLog
	for rows.Next() {
		log := ExecutionLog{}
		err := rows.Scan(
			&log.ID,
			&log.ExecutionID,
			&log.NodeID,
			&log.ActivityName,
			&log.Level,
			&log.Message,
			&log.Metadata,
			&log.LoggedAt,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan log: %w", err)
		}
		logs = append(logs, log)
	}

	return logs, rows.Err()
}

// SaveActivityTrace saves per-activity execution trace
func (db *DB) SaveActivityTrace(ctx context.Context, trace *ActivityTrace) error {
	query := `
		INSERT INTO activity_traces (execution_id, node_id, activity_name, parameters, result, started_at, completed_at, duration_ms, attempt, retry_reason, status, error_message)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
		ON CONFLICT (id) DO UPDATE SET
			status = $11,
			result = $5,
			completed_at = $7,
			duration_ms = $8,
			error_message = $12
	`

	_, err := db.conn.ExecContext(ctx, query,
		trace.ExecutionID,
		trace.NodeID,
		trace.ActivityName,
		trace.Parameters,
		trace.Result,
		trace.StartedAt,
		trace.CompletedAt,
		trace.DurationMs,
		trace.Attempt,
		trace.RetryReason,
		trace.Status,
		trace.ErrorMessage,
	)

	return err
}

// FetchActivityTraces retrieves all activity traces for an execution
func (db *DB) FetchActivityTraces(ctx context.Context, executionID string) ([]ActivityTrace, error) {
	query := `
		SELECT id, execution_id, node_id, activity_name, parameters, result, started_at, completed_at, duration_ms, attempt, retry_reason, status, error_message
		FROM activity_traces
		WHERE execution_id = $1
		ORDER BY started_at ASC
	`

	rows, err := db.conn.QueryContext(ctx, query, executionID)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch activity traces: %w", err)
	}
	defer rows.Close()

	var traces []ActivityTrace
	for rows.Next() {
		trace := ActivityTrace{}
		err := rows.Scan(
			&trace.ID,
			&trace.ExecutionID,
			&trace.NodeID,
			&trace.ActivityName,
			&trace.Parameters,
			&trace.Result,
			&trace.StartedAt,
			&trace.CompletedAt,
			&trace.DurationMs,
			&trace.Attempt,
			&trace.RetryReason,
			&trace.Status,
			&trace.ErrorMessage,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to scan trace: %w", err)
		}
		traces = append(traces, trace)
	}

	return traces, rows.Err()
}