package openapi

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

// SessionHooks provides hooks into the session lifecycle.
type SessionHooks struct {
	BeforeMessage  func(ctx context.Context, s *Session, msg *Message) error
	AfterMessage   func(ctx context.Context, s *Session, msg *Message) error
	BeforeToolCall func(ctx context.Context, s *Session, tc *ToolCall) error
	AfterToolCall  func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error
	OnReasoning    func(ctx context.Context, s *Session, reasoning string) error
	OnDone         func(ctx context.Context, s *Session) error
	OnError        func(ctx context.Context, s *Session, err error)
}

// Session manages a conversation history and tool execution logic.
type Session struct {
	client   *Client
	model    string
	messages []Message
	registry ToolRegistry
	hooks    SessionHooks
	
	// TurnTimeout is the timeout for a single LLM request.
	TurnTimeout time.Duration

	// DefaultRequestOptions allows setting default parameters for all requests in the session.
	DefaultRequestOptions ChatCompletionRequest
}

// NewSession creates a new session with the given client and model.
func NewSession(client *Client, model string) *Session {
	return &Session{
		client:      client,
		model:       model,
		registry:    NewMapRegistry(),
		TurnTimeout: 2 * time.Minute,
		DefaultRequestOptions: ChatCompletionRequest{
			CachePrompt: true,
		},
	}
}

// SetRegistry replaces the tool registry.
func (s *Session) SetRegistry(r ToolRegistry) {
	s.registry = r
}

// RegisterTool is a helper for the default MapRegistry.
func (s *Session) RegisterTool(name, description string, parameters any, handler ToolHandler) error {
	if mr, ok := s.registry.(*MapRegistry); ok {
		return mr.Register(name, description, parameters, handler)
	}
	return fmt.Errorf("current registry does not support direct registration")
}

// AddMessage adds a message to the conversation history, triggering hooks.
func (s *Session) AddMessage(ctx context.Context, msg Message) error {
	if s.hooks.BeforeMessage != nil {
		if err := s.hooks.BeforeMessage(ctx, s, &msg); err != nil {
			return err
		}
	}
	s.messages = append(s.messages, msg)
	if s.hooks.AfterMessage != nil {
		return s.hooks.AfterMessage(ctx, s, &msg)
	}
	return nil
}

// SetSystemPrompt sets the initial system prompt.
func (s *Session) SetSystemPrompt(prompt string) {
	if len(s.messages) > 0 && s.messages[0].Role == RoleSystem {
		s.messages[0].Content = prompt
	} else {
		s.messages = append([]Message{{Role: RoleSystem, Content: prompt}}, s.messages...)
	}
}

// Chat sends a user prompt and handles the response loop asynchronously.
func (s *Session) Chat(ctx context.Context, userPrompt string) (<-chan ChatCompletionChunk, <-chan error) {
	chunks := make(chan ChatCompletionChunk)
	errs := make(chan error, 1)
	go s.runLoop(ctx, userPrompt, chunks, errs)
	return chunks, errs
}

func (s *Session) runLoop(ctx context.Context, userPrompt string, chunks chan<- ChatCompletionChunk, errs chan<- error) {
	defer close(chunks)
	defer close(errs)

	if userPrompt != "" {
		if err := s.AddMessage(ctx, Message{Role: RoleUser, Content: userPrompt}); err != nil {
			s.handleError(ctx, err, errs)
			return
		}
	}

	for {
		resp, err := s.nextCompletion(ctx)
		if err != nil {
			s.handleError(ctx, err, errs)
			return
		}

		choice := resp.Choices[0]
		if err := s.processResponse(ctx, choice.Message); err != nil {
			s.handleError(ctx, err, errs)
			return
		}

		if len(choice.Message.ToolCalls) > 0 {
			if err := s.handleToolCalls(ctx, choice.Message.ToolCalls); err != nil {
				s.handleError(ctx, err, errs)
				return
			}
			continue
		}

		if s.hooks.OnDone != nil {
			s.hooks.OnDone(ctx, s)
		}
		chunks <- s.asChunk(resp, choice)
		return
	}
}

func (s *Session) nextCompletion(ctx context.Context) (*ChatCompletionResponse, error) {
	turnCtx, cancel := s.extendContext(ctx)
	defer cancel()

	req := s.DefaultRequestOptions
	req.Model = s.model
	req.Messages = s.messages
	req.Tools = s.registry.Tools(turnCtx)
	return s.client.CreateChatCompletion(turnCtx, req)
}

func (s *Session) processResponse(ctx context.Context, msg Message) error {
	if msg.ReasoningContent != "" && s.hooks.OnReasoning != nil {
		if err := s.hooks.OnReasoning(ctx, s, msg.ReasoningContent); err != nil {
			return err
		}
	}
	return s.AddMessage(ctx, msg)
}

func (s *Session) handleToolCalls(ctx context.Context, toolCalls []ToolCall) error {
	for i := range toolCalls {
		tc := &toolCalls[i]
		if s.hooks.BeforeToolCall != nil {
			if err := s.hooks.BeforeToolCall(ctx, s, tc); err != nil {
				return err
			}
		}

		// Use extended context for tool execution to prevent timeouts on long tasks
		toolCtx, cancel := s.extendContext(ctx)
		result, err := s.registry.Execute(toolCtx, s, tc.Function.Name, json.RawMessage(tc.Function.Arguments))
		cancel()
		
		if s.hooks.AfterToolCall != nil {
			if hErr := s.hooks.AfterToolCall(ctx, s, tc, result, err); hErr != nil {
				return hErr
			}
		}

		content := ""
		if err != nil {
			content = fmt.Sprintf("Error: %v", err)
		} else {
			resJSON, _ := json.Marshal(result)
			content = string(resJSON)
		}

		if err := s.AddMessage(ctx, Message{
			Role:       RoleTool,
			Name:       tc.Function.Name,
			ToolCallID: tc.ID,
			Content:    content,
		}); err != nil {
			return err
		}
	}
	return nil
}

// extendContext creates a context that resets the timeout for a new "turn",
// ignoring the parent's deadline but respecting its cancellation.
func (s *Session) extendContext(ctx context.Context) (context.Context, context.CancelFunc) {
	extCtx, cancel := context.WithTimeout(context.Background(), s.TurnTimeout)
	
	stop := make(chan struct{})
	go func() {
		select {
		case <-ctx.Done():
			cancel()
		case <-stop:
		case <-extCtx.Done():
		}
	}()

	return extCtx, func() {
		close(stop)
		cancel()
	}
}

func (s *Session) asChunk(resp *ChatCompletionResponse, choice Choice) ChatCompletionChunk {
	return ChatCompletionChunk{
		ID:      resp.ID,
		Object:  "chat.completion.chunk",
		Created: resp.Created,
		Model:   resp.Model,
		Choices: []struct {
			Index        int     `json:"index"`
			Delta        Message `json:"delta"`
			FinishReason string  `json:"finish_reason"`
		}{
			{
				Index: choice.Index,
				Delta: Message{
					Role:             choice.Message.Role,
					Content:          choice.Message.Content,
					ReasoningContent: choice.Message.ReasoningContent,
				},
				FinishReason: choice.FinishReason,
			},
		},
	}
}

func (s *Session) handleError(ctx context.Context, err error, errs chan<- error) {
	if s.hooks.OnError != nil {
		s.hooks.OnError(ctx, s, err)
	}
	errs <- err
}

func (s *Session) Messages() []Message { return s.messages }
func (s *Session) SetHooks(h SessionHooks) { s.hooks = h }