package knowledgeqa

import (
	"context"
	"errors"
	"time"

	knowledgeqastore "github.com/tidbcloud/AgentTeam/internal/store/knowledgeqa"
)

type conversationStore struct {
	store *knowledgeqastore.Store
}

func NewConversationService(store *knowledgeqastore.Store) ConversationService {
	return &conversationStore{store: store}
}

func (service *conversationStore) CreateSession(ctx context.Context, params CreateSessionParams) (Session, error) {
	session, err := service.store.CreateSession(ctx, knowledgeqastore.CreateSessionParams{
		ID:                        params.ID,
		WorkspaceID:               params.WorkspaceID,
		OrganizationID:            params.OrganizationID,
		OwnerUserID:               params.OwnerUserID,
		OwnerMemberID:             params.OwnerMemberID,
		RootMessageID:             params.RootMessageID,
		ThreadChannelID:           params.ThreadChannelID,
		RequestedKnowledgeBaseIDs: params.RequestedKnowledgeBaseIDs,
		RequestedSourceIDs:        params.RequestedSourceIDs,
		RequestedModel:            params.RequestedModel,
		RequestedReasonLevel:      params.RequestedReasonLevel,
		WebSearchEnabled:          params.WebSearchEnabled,
		Now:                       params.Now,
	})
	return sessionFromStore(session), conversationStoreError(err)
}

func (service *conversationStore) LoadSessionForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	sessionID string,
) (Session, error) {
	session, err := service.store.LoadSessionForOwner(ctx, workspaceID, userID, sessionID)
	return sessionFromStore(session), conversationStoreError(err)
}

func (service *conversationStore) ListTurnsForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	threadID string,
	limit int,
) ([]Turn, error) {
	turns, err := service.store.ListTurnsForOwner(ctx, workspaceID, userID, threadID, limit)
	if err != nil {
		return nil, conversationStoreError(err)
	}
	result := make([]Turn, 0, len(turns))
	for _, turn := range turns {
		result = append(result, turnFromStore(turn))
	}
	return result, nil
}

func (service *conversationStore) LoadCitationForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	sessionID string,
	citationID string,
) (Citation, error) {
	citation, err := service.store.LoadCitationForOwner(ctx, workspaceID, userID, sessionID, citationID)
	return citationFromStore(citation), conversationStoreError(err)
}

func (service *conversationStore) LoadFeedbackSnapshotForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	sessionID string,
	turnID string,
) (FeedbackSnapshot, error) {
	snapshot, err := service.store.LoadFeedbackSnapshotForOwner(
		ctx, workspaceID, userID, sessionID, turnID,
	)
	if err != nil {
		return FeedbackSnapshot{}, conversationStoreError(err)
	}
	citations := make([]Citation, 0, len(snapshot.Citations))
	for _, citation := range snapshot.Citations {
		citations = append(citations, citationFromStore(citation))
	}
	return FeedbackSnapshot{
		Turn:            turnFromStore(snapshot.Turn),
		ThreadChannelID: snapshot.ThreadChannelID,
		Answer:          snapshot.Answer,
		Activities:      snapshot.Activities,
		Citations:       citations,
	}, nil
}

func (service *conversationStore) CreateTurn(ctx context.Context, params CreateTurnParams) (Turn, error) {
	turn, err := service.store.CreateTurn(ctx, knowledgeqastore.CreateTurnParams{
		WorkspaceID:               params.WorkspaceID,
		OrganizationID:            params.OrganizationID,
		SessionID:                 params.SessionID,
		RequesterUserID:           params.RequesterUserID,
		UserMessageID:             params.UserMessageID,
		Question:                  params.Question,
		IdempotencyKey:            params.IdempotencyKey,
		RequestedKnowledgeBaseIDs: params.RequestedKnowledgeBaseIDs,
		RequestedSourceIDs:        params.RequestedSourceIDs,
		Now:                       params.Now,
	})
	return turnFromStore(turn), conversationStoreError(err)
}

func (service *conversationStore) RetryTurn(ctx context.Context, params RetryTurnParams) (Turn, error) {
	turn, err := service.store.RetryTurn(ctx, knowledgeqastore.RetryTurnParams{
		WorkspaceID:     params.WorkspaceID,
		SessionID:       params.SessionID,
		TurnID:          params.TurnID,
		RequesterUserID: params.RequesterUserID,
		IdempotencyKey:  params.IdempotencyKey,
		Now:             params.Now,
	})
	return turnFromStore(turn), conversationStoreError(err)
}

func (service *conversationStore) StopTurnForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	turnID string,
	now time.Time,
) (Turn, error) {
	turn, err := service.store.StopTurnForOwner(ctx, workspaceID, userID, turnID, now)
	return turnFromStore(turn), conversationStoreError(err)
}

func (service *conversationStore) DeleteSessionConversationForOwner(
	ctx context.Context,
	workspaceID string,
	userID string,
	threadID string,
	now time.Time,
) (Session, error) {
	session, err := service.store.DeleteSessionConversationForOwner(ctx, workspaceID, userID, threadID, now)
	return sessionFromStore(session), conversationStoreError(err)
}

func conversationStoreError(err error) error {
	switch {
	case errors.Is(err, knowledgeqastore.ErrNotFound):
		return ErrConversationNotFound
	case errors.Is(err, knowledgeqastore.ErrForbidden):
		return ErrConversationForbidden
	case errors.Is(err, knowledgeqastore.ErrTurnBusy):
		return ErrTurnBusy
	case errors.Is(err, knowledgeqastore.ErrQueueFull):
		return ErrQueueFull
	case errors.Is(err, knowledgeqastore.ErrRateLimited):
		return ErrRateLimited
	case errors.Is(err, knowledgeqastore.ErrInvalidState):
		return ErrInvalidState
	default:
		return err
	}
}

func sessionFromStore(session knowledgeqastore.Session) Session {
	return Session{
		ID:                        session.ID,
		WorkspaceID:               session.WorkspaceID,
		OrganizationID:            session.OrganizationID,
		OwnerUserID:               session.OwnerUserID,
		OwnerMemberID:             session.OwnerMemberID,
		RootMessageID:             session.RootMessageID,
		ThreadChannelID:           session.ThreadChannelID,
		RequestedKnowledgeBaseIDs: session.RequestedKnowledgeBaseIDs,
		RequestedSourceIDs:        session.RequestedSourceIDs,
		RequestedModel:            session.RequestedModel,
		RequestedReasonLevel:      session.RequestedReasonLevel,
		WebSearchEnabled:          session.WebSearchEnabled,
		ActiveTurnID:              session.ActiveTurnID,
		Status:                    session.Status,
	}
}

func turnFromStore(turn knowledgeqastore.Turn) Turn {
	debugModelOutput := ""
	if turn.Status == "failed" {
		debugModelOutput = turn.DebugModelOutput
	}
	return Turn{
		ID:                        turn.ID,
		WorkspaceID:               turn.WorkspaceID,
		OrganizationID:            turn.OrganizationID,
		SessionID:                 turn.SessionID,
		RequesterUserID:           turn.RequesterUserID,
		UserMessageID:             turn.UserMessageID,
		AssistantMessageID:        turn.AssistantMessageID,
		Question:                  turn.Question,
		RequestedKnowledgeBaseIDs: append([]string(nil), turn.RequestedKnowledgeBaseIDs...),
		RequestedSourceIDs:        append([]string(nil), turn.RequestedSourceIDs...),
		AuthorizedRoots:           authorizedRootsFromStore(turn.AuthorizedRoots),
		RequestedModel:            turn.RequestedModel,
		RequestedReasonLevel:      turn.RequestedReasonLevel,
		EffectiveModel:            turn.EffectiveModel,
		EffectiveReasonLevel:      turn.EffectiveReasonLevel,
		OrganizationPolicyVersion: turn.OrganizationPolicyVersion,
		RuntimeVersion:            turn.RuntimeVersion,
		PromptVersion:             turn.PromptVersion,
		CurrentAttemptID:          turn.CurrentAttemptID,
		Status:                    turn.Status,
		ErrorCode:                 turn.ErrorCode,
		DebugModelOutput:          debugModelOutput,
		Activities:                turn.Activities,

		WebSearchEnabled: turn.WebSearchEnabled,
	}
}

func authorizedRootsFromStore(roots []knowledgeqastore.AuthorizedRoot) []AuthorizedRoot {
	result := make([]AuthorizedRoot, 0, len(roots))
	for _, root := range roots {
		result = append(result, AuthorizedRoot{
			KnowledgeBaseID: root.KnowledgeBaseID,
			SourceID:        root.SourceID,
			RootKind:        root.RootKind,
			RepoRole:        root.RepoRole,
			Revision:        root.Revision,
			FreshnessStatus: root.FreshnessStatus,
		})
	}
	return result
}

func citationFromStore(citation knowledgeqastore.Citation) Citation {
	return Citation{
		ID:              citation.ID,
		TurnID:          citation.TurnID,
		CitationKind:    citation.CitationKind,
		KnowledgeBaseID: citation.KnowledgeBaseID,
		SourceID:        citation.SourceID,
		RootKind:        citation.RootKind,
		RepoRole:        citation.RepoRole,
		Revision:        citation.Revision,
		ResolvedRef:     citation.ResolvedRef,
		ResolvedRefKind: citation.ResolvedRefKind,
		Path:            citation.Path,
		LineStart:       citation.LineStart,
		LineEnd:         citation.LineEnd,
		URL:             citation.URL,
		Title:           citation.Title,
		Snippet:         citation.Snippet,
		FetchedAt:       citation.FetchedAt,
	}
}