package thinking

import (
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/xlgmokha/mcp/pkg/mcp"
)

// ThinkingOperations provides thinking session management operations
type ThinkingOperations struct {
	sessions     map[string]*ThinkingSession
	branches     map[string]*Branch
	mu           sync.RWMutex
	persistFile  string
}

// ThinkingSession represents a thinking session with sequential thoughts
type ThinkingSession struct {
	ID             string    `json:"id"`
	Thoughts       []Thought `json:"thoughts"`
	CurrentThought int       `json:"current_thought"`
	TotalThoughts  int       `json:"total_thoughts"`
	Status         string    `json:"status"` // "active", "completed"
	CreatedAt      time.Time `json:"created_at"`
	LastActivity   time.Time `json:"last_activity"`
	ActiveBranches []string  `json:"active_branches"`
}

// Branch represents a reasoning branch with its own thought sequence
type Branch struct {
	ID             string    `json:"id"`
	SessionID      string    `json:"session_id"`
	FromThought    int       `json:"from_thought"`
	Thoughts       []Thought `json:"thoughts"`
	CreatedAt      time.Time `json:"created_at"`
	LastActivity   time.Time `json:"last_activity"`
}

// Thought represents a single thought in the sequence
type Thought struct {
	Number            int    `json:"number"`
	Content           string `json:"content"`
	IsRevision        bool   `json:"is_revision,omitempty"`
	RevisesThought    *int   `json:"revises_thought,omitempty"`
	BranchFromThought *int   `json:"branch_from_thought,omitempty"`
	BranchID          string `json:"branch_id,omitempty"`
	NeedsMoreThoughts bool   `json:"needs_more_thoughts,omitempty"`
}

// ThinkingResponse represents the response structure
type ThinkingResponse struct {
	Thought            string   `json:"thought"`
	ThoughtNumber      int      `json:"thought_number"`
	TotalThoughts      int      `json:"total_thoughts"`
	NextThoughtNeeded  bool     `json:"next_thought_needed"`
	Status             string   `json:"status"`
	Solution           string   `json:"solution,omitempty"`
	SessionID          string   `json:"session_id"`
	ThoughtHistorySize int      `json:"thought_history_size"`
	ActiveBranches     []string `json:"active_branches"`
	BranchContext      string   `json:"branch_context,omitempty"`
}

// NewThinkingOperations creates a new ThinkingOperations helper
func NewThinkingOperations(persistFile string) *ThinkingOperations {
	thinking := &ThinkingOperations{
		sessions:    make(map[string]*ThinkingSession),
		branches:    make(map[string]*Branch),
		persistFile: persistFile,
	}

	// Load existing sessions if persistence file is provided
	if persistFile != "" {
		thinking.loadSessions()
	}

	return thinking
}

// New creates a new Sequential Thinking MCP server
func New() *mcp.Server {
	return NewWithPersistence("")
}

// NewWithPersistence creates a new Sequential Thinking MCP server with optional persistence
func NewWithPersistence(persistFile string) *mcp.Server {
	thinking := NewThinkingOperations(persistFile)
	builder := mcp.NewServerBuilder("mcp-sequential-thinking", "1.0.0")

	// Add sequentialthinking tool
	builder.AddTool(mcp.NewTool("sequentialthinking", "A detailed tool for dynamic and reflective problem-solving through thoughts. This tool helps analyze problems through a flexible thinking process that can adapt and evolve. Each thought can build on, question, or revise previous insights as understanding deepens.", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"thought": map[string]interface{}{
				"type":        "string",
				"description": "Your current thinking step",
			},
			"nextThoughtNeeded": map[string]interface{}{
				"type":        "boolean",
				"description": "Whether another thought step is needed",
			},
			"thoughtNumber": map[string]interface{}{
				"type":        "integer",
				"minimum":     1,
				"description": "Current thought number",
			},
			"totalThoughts": map[string]interface{}{
				"type":        "integer",
				"minimum":     1,
				"description": "Estimated total thoughts needed",
			},
			"isRevision": map[string]interface{}{
				"type":        "boolean",
				"description": "Whether this revises previous thinking",
				"default":     false,
			},
			"revisesThought": map[string]interface{}{
				"type":        "integer",
				"minimum":     1,
				"description": "Which thought is being reconsidered",
			},
			"branchFromThought": map[string]interface{}{
				"type":        "integer",
				"minimum":     1,
				"description": "Branching point thought number",
			},
			"branchId": map[string]interface{}{
				"type":        "string",
				"description": "Branch identifier",
			},
			"needsMoreThoughts": map[string]interface{}{
				"type":        "boolean",
				"description": "If more thoughts are needed",
				"default":     false,
			},
			"sessionId": map[string]interface{}{
				"type":        "string",
				"description": "Session ID for thought continuity (optional, auto-generated if not provided)",
			},
		},
		"required": []string{"thought", "nextThoughtNeeded", "thoughtNumber", "totalThoughts"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		return thinking.handleSequentialThinking(req)
	}))

	// Add get_session_history tool
	builder.AddTool(mcp.NewTool("get_session_history", "Get the complete thought history for a thinking session", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"sessionId": map[string]interface{}{
				"type":        "string",
				"description": "Session ID to get history for",
			},
		},
		"required": []string{"sessionId"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		return thinking.handleGetSessionHistory(req)
	}))

	// Add list_sessions tool
	builder.AddTool(mcp.NewTool("list_sessions", "List all active thinking sessions", map[string]interface{}{
		"type":       "object",
		"properties": map[string]interface{}{},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		return thinking.handleListSessions(req)
	}))

	// Add get_branch_history tool
	builder.AddTool(mcp.NewTool("get_branch_history", "Get the thought history for a specific reasoning branch", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"branchId": map[string]interface{}{
				"type":        "string",
				"description": "Branch ID to get history for",
			},
		},
		"required": []string{"branchId"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		return thinking.handleGetBranchHistory(req)
	}))

	// Add clear_session tool
	builder.AddTool(mcp.NewTool("clear_session", "Clear a thinking session and all its branches", map[string]interface{}{
		"type": "object",
		"properties": map[string]interface{}{
			"sessionId": map[string]interface{}{
				"type":        "string",
				"description": "Session ID to clear",
			},
		},
		"required": []string{"sessionId"},
	}, func(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
		return thinking.handleClearSession(req)
	}))

	return builder.Build()
}

// Helper methods for ThinkingOperations

func (thinking *ThinkingOperations) handleSequentialThinking(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
	// Parse and validate input parameters
	params, err := thinking.parseThinkingParameters(req.Arguments)
	if err != nil {
		return mcp.NewToolError(err.Error()), nil
	}

	thinking.mu.Lock()
	defer thinking.mu.Unlock()

	// Get or create session
	session := thinking.getOrCreateSession(params.SessionID)
	
	// Create thought object
	currentThought := Thought{
		Number:            params.ThoughtNumber,
		Content:           params.Thought,
		IsRevision:        params.IsRevision,
		RevisesThought:    params.RevisesThought,
		BranchFromThought: params.BranchFromThought,
		BranchID:          params.BranchID,
		NeedsMoreThoughts: params.NeedsMoreThoughts,
	}

	// Handle branching
	var activeBranch *Branch
	if params.BranchFromThought != nil && params.BranchID != "" {
		activeBranch = thinking.getOrCreateBranch(session.ID, params.BranchID, *params.BranchFromThought)
		activeBranch.Thoughts = append(activeBranch.Thoughts, currentThought)
		activeBranch.LastActivity = time.Now()
	} else {
		// Add to main session
		session.Thoughts = append(session.Thoughts, currentThought)
	}

	// Update session state
	session.CurrentThought = params.ThoughtNumber
	session.TotalThoughts = params.TotalThoughts
	session.LastActivity = time.Now()

	// Save to persistence if configured
	thinking.saveSessions()

	// Determine status
	status := "thinking"
	if !params.NextThoughtNeeded {
		status = "completed"
	} else if params.ThoughtNumber >= params.TotalThoughts && !params.NeedsMoreThoughts {
		status = "completed"
	}
	session.Status = status

	// Create response with session context
	response := ThinkingResponse{
		Thought:            params.Thought,
		ThoughtNumber:      params.ThoughtNumber,
		TotalThoughts:      params.TotalThoughts,
		NextThoughtNeeded:  params.NextThoughtNeeded,
		Status:             status,
		SessionID:          session.ID,
		ThoughtHistorySize: len(session.Thoughts),
		ActiveBranches:     session.ActiveBranches,
	}

	// Add branch context if applicable
	if activeBranch != nil {
		response.BranchContext = fmt.Sprintf("Branch %s (from thought %d)", params.BranchID, *params.BranchFromThought)
	}

	// If this is the final thought, try to extract a solution
	if status == "completed" {
		response.Solution = thinking.extractSolution(params.Thought)
	}

	// Format the result with session context
	resultText := thinking.formatThinkingResultWithSession(response, currentThought, session, activeBranch)

	return mcp.NewToolResult(mcp.NewTextContent(resultText)), nil
}

// Helper methods

// PersistentData represents the data structure for persistence
type PersistentData struct {
	Sessions map[string]*ThinkingSession `json:"sessions"`
	Branches map[string]*Branch          `json:"branches"`
}

// loadSessions loads sessions from persistence file
func (thinking *ThinkingOperations) loadSessions() error {
	if thinking.persistFile == "" {
		return nil
	}

	if _, err := os.Stat(thinking.persistFile); os.IsNotExist(err) {
		return nil // File doesn't exist, start fresh
	}

	data, err := os.ReadFile(thinking.persistFile)
	if err != nil {
		return err
	}

	if len(data) == 0 {
		return nil // Empty file
	}

	var persistentData PersistentData
	if err := json.Unmarshal(data, &persistentData); err != nil {
		return err
	}

	thinking.sessions = persistentData.Sessions
	thinking.branches = persistentData.Branches

	// Initialize maps if nil
	if thinking.sessions == nil {
		thinking.sessions = make(map[string]*ThinkingSession)
	}
	if thinking.branches == nil {
		thinking.branches = make(map[string]*Branch)
	}

	return nil
}

// saveSessions saves sessions to persistence file
func (thinking *ThinkingOperations) saveSessions() error {
	if thinking.persistFile == "" {
		return nil
	}

	persistentData := PersistentData{
		Sessions: thinking.sessions,
		Branches: thinking.branches,
	}

	data, err := json.MarshalIndent(persistentData, "", "  ")
	if err != nil {
		return err
	}

	return os.WriteFile(thinking.persistFile, data, 0644)
}

// ThinkingParameters holds parsed parameters for thinking operations
type ThinkingParameters struct {
	Thought           string
	NextThoughtNeeded bool
	ThoughtNumber     int
	TotalThoughts     int
	IsRevision        bool
	RevisesThought    *int
	BranchFromThought *int
	BranchID          string
	NeedsMoreThoughts bool
	SessionID         string
}

// parseThinkingParameters parses and validates input parameters
func (thinking *ThinkingOperations) parseThinkingParameters(args map[string]interface{}) (*ThinkingParameters, error) {
	params := &ThinkingParameters{}

	// Required parameters
	thought, ok := args["thought"].(string)
	if !ok {
		return nil, fmt.Errorf("thought parameter is required and must be a string")
	}
	params.Thought = thought

	nextThoughtNeeded, ok := args["nextThoughtNeeded"].(bool)
	if !ok {
		return nil, fmt.Errorf("nextThoughtNeeded parameter is required and must be a boolean")
	}
	params.NextThoughtNeeded = nextThoughtNeeded

	if tn, ok := args["thoughtNumber"]; ok {
		switch v := tn.(type) {
		case float64:
			params.ThoughtNumber = int(v)
		case int:
			params.ThoughtNumber = v
		default:
			return nil, fmt.Errorf("thoughtNumber must be a number")
		}
		if params.ThoughtNumber < 1 {
			return nil, fmt.Errorf("thoughtNumber must be >= 1")
		}
	} else {
		params.ThoughtNumber = 1
	}

	if tt, ok := args["totalThoughts"]; ok {
		switch v := tt.(type) {
		case float64:
			params.TotalThoughts = int(v)
		case int:
			params.TotalThoughts = v
		default:
			return nil, fmt.Errorf("totalThoughts must be a number")
		}
		if params.TotalThoughts < 1 {
			return nil, fmt.Errorf("totalThoughts must be >= 1")
		}
	} else {
		params.TotalThoughts = 1
	}

	// Optional parameters
	if ir, ok := args["isRevision"].(bool); ok {
		params.IsRevision = ir
	}

	if rt, ok := args["revisesThought"]; ok && params.IsRevision {
		switch v := rt.(type) {
		case float64:
			val := int(v)
			params.RevisesThought = &val
		case int:
			params.RevisesThought = &v
		default:
			return nil, fmt.Errorf("revisesThought must be a number")
		}
	}

	if bft, ok := args["branchFromThought"]; ok {
		switch v := bft.(type) {
		case float64:
			val := int(v)
			params.BranchFromThought = &val
		case int:
			params.BranchFromThought = &v
		default:
			return nil, fmt.Errorf("branchFromThought must be a number")
		}
	}

	if bid, ok := args["branchId"].(string); ok {
		params.BranchID = bid
	}

	if nmt, ok := args["needsMoreThoughts"].(bool); ok {
		params.NeedsMoreThoughts = nmt
	}

	if sid, ok := args["sessionId"].(string); ok {
		params.SessionID = sid
	}

	return params, nil
}

// getOrCreateSession gets existing session or creates new one
func (thinking *ThinkingOperations) getOrCreateSession(sessionID string) *ThinkingSession {
	if sessionID == "" {
		sessionID = thinking.generateSessionID()
	}

	if session, exists := thinking.sessions[sessionID]; exists {
		return session
	}

	session := &ThinkingSession{
		ID:             sessionID,
		Thoughts:       make([]Thought, 0),
		CurrentThought: 0,
		TotalThoughts:  1,
		Status:         "active",
		CreatedAt:      time.Now(),
		LastActivity:   time.Now(),
		ActiveBranches: make([]string, 0),
	}

	thinking.sessions[sessionID] = session
	return session
}

// getOrCreateBranch gets existing branch or creates new one
func (thinking *ThinkingOperations) getOrCreateBranch(sessionID, branchID string, fromThought int) *Branch {
	if branch, exists := thinking.branches[branchID]; exists {
		return branch
	}

	branch := &Branch{
		ID:           branchID,
		SessionID:    sessionID,
		FromThought:  fromThought,
		Thoughts:     make([]Thought, 0),
		CreatedAt:    time.Now(),
		LastActivity: time.Now(),
	}

	thinking.branches[branchID] = branch

	// Add to session's active branches
	if session, exists := thinking.sessions[sessionID]; exists {
		session.ActiveBranches = append(session.ActiveBranches, branchID)
	}

	return branch
}

// generateSessionID generates a unique session ID
func (thinking *ThinkingOperations) generateSessionID() string {
	return fmt.Sprintf("session_%d", time.Now().UnixNano())
}

// formatThinkingResultWithSession formats result with session context
func (thinking *ThinkingOperations) formatThinkingResultWithSession(response ThinkingResponse, thought Thought, session *ThinkingSession, branch *Branch) string {
	var result strings.Builder

	// Header with session info
	result.WriteString(fmt.Sprintf("🧠 Sequential Thinking - Thought %d/%d (Session: %s)\n",
		response.ThoughtNumber, response.TotalThoughts, response.SessionID))
	result.WriteString("═══════════════════════════════════════\n\n")

	// Session context
	result.WriteString(fmt.Sprintf("📋 Session Info:\n"))
	result.WriteString(fmt.Sprintf("   • Session ID: %s\n", session.ID))
	result.WriteString(fmt.Sprintf("   • Total thoughts in session: %d\n", len(session.Thoughts)))
	result.WriteString(fmt.Sprintf("   • Active branches: %d\n", len(session.ActiveBranches)))
	if response.BranchContext != "" {
		result.WriteString(fmt.Sprintf("   • Branch context: %s\n", response.BranchContext))
	}
	result.WriteString("\n")

	// Context information
	var contextInfo []string
	if thought.IsRevision && thought.RevisesThought != nil {
		contextInfo = append(contextInfo, fmt.Sprintf("Revising thought %d", *thought.RevisesThought))
	}
	if thought.BranchFromThought != nil {
		contextInfo = append(contextInfo, fmt.Sprintf("Branching from thought %d", *thought.BranchFromThought))
		if thought.BranchID != "" {
			contextInfo = append(contextInfo, fmt.Sprintf("Branch: %s", thought.BranchID))
		}
	}
	if thought.NeedsMoreThoughts {
		contextInfo = append(contextInfo, "Requesting additional thoughts beyond initial estimate")
	}

	if len(contextInfo) > 0 {
		result.WriteString("📝 Context: " + strings.Join(contextInfo, ", ") + "\n\n")
	}

	// Main thought content
	result.WriteString("💭 Current Thought:\n")
	result.WriteString(response.Thought + "\n\n")

	// Progress indicator
	progressBar := thinking.createProgressBar(response.ThoughtNumber, response.TotalThoughts)
	result.WriteString("📊 Progress: " + progressBar + "\n\n")

	// Status
	statusEmoji := "🔄"
	if response.Status == "completed" {
		statusEmoji = "✅"
	}
	result.WriteString(fmt.Sprintf("%s Status: %s\n", statusEmoji, response.Status))

	if response.NextThoughtNeeded {
		result.WriteString("⏭️  Next thought needed\n")
	} else {
		result.WriteString("🏁 Thinking sequence complete\n")
	}

	// Solution if available
	if response.Solution != "" {
		result.WriteString("\n🎯 Extracted Solution:\n")
		result.WriteString(response.Solution + "\n")
	}

	// JSON data for programmatic access
	result.WriteString("\n📋 Structured Data:\n")
	jsonData, _ := json.MarshalIndent(response, "", "  ")
	result.WriteString("```json\n")
	result.WriteString(string(jsonData))
	result.WriteString("\n```")

	return result.String()
}

func (thinking *ThinkingOperations) extractSolution(finalThought string) string {
	// Simple heuristic to extract a solution from the final thought
	content := strings.ToLower(finalThought)

	// Look for solution indicators (with and without colons)
	solutionKeywords := []string{
		"solution:",
		"solution is",
		"answer:",
		"answer is",
		"conclusion:",
		"final answer:",
		"result:",
		"therefore:",
		"therefore,",
		"therefore the",
		"in conclusion:",
	}

	for _, keyword := range solutionKeywords {
		if idx := strings.Index(content, keyword); idx != -1 {
			// Extract text after the keyword
			remaining := strings.TrimSpace(finalThought[idx+len(keyword):])
			if len(remaining) > 0 {
				// Take up to the first sentence or 200 characters
				if sentences := strings.Split(remaining, "."); len(sentences) > 0 {
					solution := strings.TrimSpace(sentences[0])
					if len(solution) > 200 {
						solution = solution[:200] + "..."
					}
					if solution != "" {
						return solution
					}
				}
			}
		}
	}

	// If no explicit solution found, return the last sentence or a portion
	sentences := strings.Split(strings.TrimSpace(finalThought), ".")
	if len(sentences) > 0 {
		lastSentence := strings.TrimSpace(sentences[len(sentences)-1])
		if len(lastSentence) > 200 {
			lastSentence = lastSentence[:200] + "..."
		}
		if lastSentence != "" {
			return lastSentence
		}
	}

	return "Solution extracted from final thought"
}

func (thinking *ThinkingOperations) formatThinkingResult(response ThinkingResponse, thought Thought, contextInfo []string) string {
	var result strings.Builder

	// Header
	result.WriteString(fmt.Sprintf("🧠 Sequential Thinking - Thought %d/%d\n",
		response.ThoughtNumber, response.TotalThoughts))
	result.WriteString("═══════════════════════════════════════\n\n")

	// Context information if any
	if len(contextInfo) > 0 {
		result.WriteString("📝 Context: " + strings.Join(contextInfo, ", ") + "\n\n")
	}

	// Main thought content
	result.WriteString("💭 Current Thought:\n")
	result.WriteString(response.Thought + "\n\n")

	// Progress indicator
	progressBar := thinking.createProgressBar(response.ThoughtNumber, response.TotalThoughts)
	result.WriteString("📊 Progress: " + progressBar + "\n\n")

	// Status
	statusEmoji := "🔄"
	if response.Status == "completed" {
		statusEmoji = "✅"
	}
	result.WriteString(fmt.Sprintf("%s Status: %s\n", statusEmoji, response.Status))

	if response.NextThoughtNeeded {
		result.WriteString("⏭️  Next thought needed\n")
	} else {
		result.WriteString("🏁 Thinking sequence complete\n")
	}

	// Solution if available
	if response.Solution != "" {
		result.WriteString("\n🎯 Extracted Solution:\n")
		result.WriteString(response.Solution + "\n")
	}

	// JSON data for programmatic access
	result.WriteString("\n📋 Structured Data:\n")
	jsonData, _ := json.MarshalIndent(response, "", "  ")
	result.WriteString("```json\n")
	result.WriteString(string(jsonData))
	result.WriteString("\n```")

	return result.String()
}

func (thinking *ThinkingOperations) createProgressBar(current, total int) string {
	if total <= 0 {
		return "[████████████████████] 100%"
	}

	percentage := float64(current) / float64(total) * 100
	if percentage > 100 {
		percentage = 100
	}

	barLength := 20
	filledLength := int(percentage / 100 * float64(barLength))

	bar := "["
	for i := 0; i < barLength; i++ {
		if i < filledLength {
			bar += "█"
		} else {
			bar += "░"
		}
	}
	bar += "] " + strconv.Itoa(int(percentage)) + "%"

	return bar
}

// New tool handlers for session management

func (thinking *ThinkingOperations) handleGetSessionHistory(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
	sessionID, ok := req.Arguments["sessionId"].(string)
	if !ok {
		return mcp.NewToolError("sessionId parameter is required"), nil
	}

	thinking.mu.RLock()
	defer thinking.mu.RUnlock()

	session, exists := thinking.sessions[sessionID]
	if !exists {
		return mcp.NewToolError(fmt.Sprintf("Session %s not found", sessionID)), nil
	}

	result := map[string]interface{}{
		"sessionId":      session.ID,
		"status":         session.Status,
		"totalThoughts":  len(session.Thoughts),
		"createdAt":      session.CreatedAt,
		"lastActivity":   session.LastActivity,
		"activeBranches": session.ActiveBranches,
		"thoughts":       session.Thoughts,
	}

	jsonData, _ := json.MarshalIndent(result, "", "  ")
	return mcp.NewToolResult(mcp.NewTextContent(string(jsonData))), nil
}

func (thinking *ThinkingOperations) handleListSessions(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
	thinking.mu.RLock()
	defer thinking.mu.RUnlock()

	sessions := make([]map[string]interface{}, 0, len(thinking.sessions))
	for _, session := range thinking.sessions {
		sessionInfo := map[string]interface{}{
			"sessionId":      session.ID,
			"status":         session.Status,
			"thoughtCount":   len(session.Thoughts),
			"branchCount":    len(session.ActiveBranches),
			"createdAt":      session.CreatedAt,
			"lastActivity":   session.LastActivity,
		}
		sessions = append(sessions, sessionInfo)
	}

	result := map[string]interface{}{
		"sessions": sessions,
		"total":    len(sessions),
	}

	jsonData, _ := json.MarshalIndent(result, "", "  ")
	return mcp.NewToolResult(mcp.NewTextContent(string(jsonData))), nil
}

func (thinking *ThinkingOperations) handleGetBranchHistory(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
	branchID, ok := req.Arguments["branchId"].(string)
	if !ok {
		return mcp.NewToolError("branchId parameter is required"), nil
	}

	thinking.mu.RLock()
	defer thinking.mu.RUnlock()

	branch, exists := thinking.branches[branchID]
	if !exists {
		return mcp.NewToolError(fmt.Sprintf("Branch %s not found", branchID)), nil
	}

	result := map[string]interface{}{
		"branchId":     branch.ID,
		"sessionId":    branch.SessionID,
		"fromThought":  branch.FromThought,
		"thoughtCount": len(branch.Thoughts),
		"createdAt":    branch.CreatedAt,
		"lastActivity": branch.LastActivity,
		"thoughts":     branch.Thoughts,
	}

	jsonData, _ := json.MarshalIndent(result, "", "  ")
	return mcp.NewToolResult(mcp.NewTextContent(string(jsonData))), nil
}

func (thinking *ThinkingOperations) handleClearSession(req mcp.CallToolRequest) (mcp.CallToolResult, error) {
	sessionID, ok := req.Arguments["sessionId"].(string)
	if !ok {
		return mcp.NewToolError("sessionId parameter is required"), nil
	}

	thinking.mu.Lock()
	defer thinking.mu.Unlock()

	session, exists := thinking.sessions[sessionID]
	if !exists {
		return mcp.NewToolError(fmt.Sprintf("Session %s not found", sessionID)), nil
	}

	// Remove associated branches
	for _, branchID := range session.ActiveBranches {
		delete(thinking.branches, branchID)
	}

	// Remove the session
	delete(thinking.sessions, sessionID)

	// Save to persistence if configured
	thinking.saveSessions()

	result := map[string]interface{}{
		"message":       fmt.Sprintf("Session %s and %d branches cleared", sessionID, len(session.ActiveBranches)),
		"sessionId":     sessionID,
		"branchesCount": len(session.ActiveBranches),
	}

	jsonData, _ := json.MarshalIndent(result, "", "  ")
	return mcp.NewToolResult(mcp.NewTextContent(string(jsonData))), nil
}