import { Inject, Logger, Provide, Scope, ScopeEnum } from '@midwayjs/core';
import { ILogger } from '@midwayjs/logger';
import { NetaClawSessionService } from '../gateway/session.js';
import { LLMMessage, LLMMessageWithId } from '../plugins/plugin_entry.js';
import type {
  SessionTreeEntry,
  SessionTreeSession,
  SessionTreeSubagentBatchMetadata,
  SessionTreeSubagentResultMetadata,
  SubagentBatchEventProjection,
  SubagentEvidenceSummaryProjection,
  SubagentProcessEventProjection,
  SubagentToolRuntimeRouteProjection,
} from '../session-tree/types.js';

export interface ChatOrchestratorRunnerParams {
  sessionId: string;
  assistantMessageId: number;
  userMessageId: number;
  userMessage: string;
  history: LLMMessage[];
  subagents: ChatOrchestratorSubagentContext;
}

export interface SessionTreeChatOrchestratorRunnerParams {
  sessionId: string;
  userEntryId: string;
  assistantEntryId: string;
  userMessage: string;
  history: LLMMessage[];
  subagents: ChatOrchestratorSubagentContext;
  onAssistantDelta?: (text: string) => void;
}

export interface SessionTreeChatRealtimeSink {
  onEntryAdded?: (event: { sessionId: string; entry: SessionTreeEntry; session?: SessionTreeSession }) => void;
  onEntryUpdated?: (event: { sessionId: string; entry: SessionTreeEntry; session?: SessionTreeSession }) => void;
  onLeafChanged?: (event: { sessionId: string; leafEntryId: string | null; session?: SessionTreeSession }) => void;
  onAssistantStreamStart?: (event: { sessionId: string; assistantEntryId: string }) => void;
  onAssistantStreamDelta?: (event: { sessionId: string; assistantEntryId?: string; text: string }) => void;
  onAssistantStreamEnd?: (event: { sessionId: string; assistantEntryId: string }) => void;
  onSubagentRunStarted?: (event: { sessionId: string; batchId: string; parentEntryId?: string | null; mode: 'single' | 'parallel'; tasks?: ChatOrchestratorSubagentTaskDescriptor[] }) => void;
  onSubagentEvent?: (event: {
    sessionId: string;
    batchId: string;
    parentEntryId?: string | null;
    sequence: number;
    timestamp: string;
    event: {
      kind: 'progress' | 'result' | 'status';
      status: 'in_flight' | 'completed' | 'failed';
      subagent: ChatOrchestratorSubagentRecord;
    };
  }) => void;
  onSubagentRunCompleted?: (event: { sessionId: string; batchId: string; parentEntryId?: string | null; results: ChatOrchestratorSubagentRecord[] }) => void;
}

export interface ChatOrchestratorRunnerResult {
  finalContent: string;
  usage?: {
    inputTokens?: number;
    outputTokens?: number;
    [key: string]: unknown;
  };
  toolCallCount?: number;
  thinking?: string;
  metadata?: Record<string, unknown>;
}

export interface ChatOrchestratorSubagentTaskDescriptor {
  goal?: string;
  context?: string;
  mode?: string;
  agentName?: string | null;
  agentId?: number | null;
  toolNames?: string[];
  [key: string]: unknown;
}

export interface ChatOrchestratorSubagentRecord {
  id?: number;
  sortOrder?: number;
  name?: string | null;
  goal?: string | null;
  status?: string;
  summary?: string | null;
  error?: string | null;
  [key: string]: unknown;
}

export interface ChatOrchestratorSubagentBatch {
  batchId: string;
  sessionId: string;
  parentMessageId: number;
  mode: 'single' | 'parallel';
  tasks: ChatOrchestratorSubagentTaskDescriptor[];
  status: 'running' | 'completed';
  results: ChatOrchestratorSubagentRecord[];
  metadata?: SessionTreeSubagentBatchMetadata;
}

export interface ChatOrchestratorSubagentBatchStartEvent {
  batchId: string;
  sessionId: string;
  parentMessageId: number;
  mode: 'single' | 'parallel';
  tasks?: ChatOrchestratorSubagentTaskDescriptor[];
  toolRuntimeRoutes?: SubagentToolRuntimeRouteProjection[];
}

export interface ChatOrchestratorSubagentUpdateEvent {
  batchId: string;
  sessionId: string;
  parentMessageId: number;
  sequence: number;
  timestamp: string;
  event: {
    kind: 'progress' | 'result' | 'status';
    status: 'in_flight' | 'completed' | 'failed';
    subagent: ChatOrchestratorSubagentRecord;
  };
}

export interface ChatOrchestratorSubagentDoneEvent {
  batchId: string;
  sessionId: string;
  parentMessageId: number;
  results: ChatOrchestratorSubagentRecord[];
}

export interface ChatOrchestratorSubagentContext {
  sessionId: string;
  parentMessageId: number;
  batches: ChatOrchestratorSubagentBatch[];
  onBatchStart: (event: ChatOrchestratorSubagentBatchStartEvent) => void;
  onUpdate: (event: ChatOrchestratorSubagentUpdateEvent) => void;
  onDone: (event: ChatOrchestratorSubagentDoneEvent) => void;
}

@Provide()
@Scope(ScopeEnum.Singleton)
export class NetaClawChatOrchestratorService {
  @Inject()
  sessionService: NetaClawSessionService;

  @Logger()
  logger?: ILogger;

  private sessionTreeSubagentParentEntryIds = new Map<string, string>();
  private sessionTreeSubagentBatchEntryIds = new Map<string, string>();

  private buildSubagentBatchMetadata(
    batch: ChatOrchestratorSubagentBatch,
    nextEvent?: SubagentBatchEventProjection,
    extra?: SessionTreeSubagentBatchMetadata,
  ): SessionTreeSubagentBatchMetadata {
    const previousEvents = Array.isArray(batch.metadata?.events)
      ? batch.metadata!.events as SubagentBatchEventProjection[]
      : [];
    const events = nextEvent ? previousEvents.concat(nextEvent).slice(-12) : previousEvents;
    const previousProcessEvents = Array.isArray(batch.metadata?.processEvents)
      ? batch.metadata!.processEvents as SubagentProcessEventProjection[]
      : [];
    const nextProcessEvents = Array.isArray(extra?.processEvents)
      ? extra.processEvents as SubagentProcessEventProjection[]
      : previousProcessEvents;

    return {
      ...(batch.metadata ?? {}),
      ...(extra ?? {}),
      latestEvent: nextEvent ?? batch.metadata?.latestEvent,
      events,
      processEvents: nextProcessEvents,
    };
  }

  private async flushSubagentBatchEntryUpdate(params: {
    sessionId: string;
    batch: ChatOrchestratorSubagentBatch;
    batchEntryId: string;
    agentId?: number;
    realtimeSink?: SessionTreeChatRealtimeSink;
  }): Promise<void> {
    const updated = await this.sessionService.updateSubagentBatchEntry(params.sessionId, params.batchEntryId, {
      batchId: params.batch.batchId,
      mode: params.batch.mode,
      tasks: params.batch.tasks.map(task => ({ ...task })),
      status: params.batch.status,
      metadata: params.batch.metadata,
      agentId: params.agentId,
    });
    params.realtimeSink?.onEntryUpdated?.({
      sessionId: params.sessionId,
      entry: updated.entry,
      session: updated.session,
    });
  }

  async executeChat(params: {
    sessionId?: string;
    requestedAgentName: string;
    message: string;
    agentId?: number;
    userId?: string;
    prepareHistory?: (params: { sessionId: string; userMessageId: number; history: LLMMessage[] }) => Promise<LLMMessage[]>;
    runner: (params: ChatOrchestratorRunnerParams) => Promise<ChatOrchestratorRunnerResult>;
  }) {
    const sessionId = await this.sessionService.getOrCreateSession(
      params.sessionId,
      params.requestedAgentName,
      params.agentId,
      params.userId
    );

    const savedUserMessage = await this.sessionService.saveMessage(sessionId, {
      role: 'user',
      content: params.message,
    });

    const initialHistory = await this.sessionService.loadHistoryWithId(sessionId, { view: 'compacted' });
    const runnerHistory = this.excludeMessageById(initialHistory, savedUserMessage.id);
    const history = params.prepareHistory
      ? await params.prepareHistory({ sessionId, userMessageId: savedUserMessage.id, history: runnerHistory })
      : runnerHistory;
    const assistantPlaceholder = await this.sessionService.saveMessage(sessionId, {
      role: 'assistant',
      content: '',
      metadata: { inFlight: true },
    });
    const subagentBatches: ChatOrchestratorSubagentBatch[] = [];
    const subagents = this.createSubagentContext(sessionId, assistantPlaceholder.id, subagentBatches);

    let result: ChatOrchestratorRunnerResult;
    try {
      result = await params.runner({
        sessionId,
        assistantMessageId: assistantPlaceholder.id,
        userMessageId: savedUserMessage.id,
        userMessage: params.message,
        history,
        subagents,
      });
    } catch (error: any) {
      const errorMessage = error?.message || 'Agent execution failed';
      try {
        await this.sessionService.updateMessage(assistantPlaceholder.id, {
          content: errorMessage,
          metadata: {
            inFlight: false,
            ...(subagentBatches.length > 0
              ? { subagentBatches: this.cloneSubagentBatches(subagentBatches) }
              : {}),
            error: {
              message: errorMessage,
              name: error?.name || 'Error',
            },
          },
        });
      } catch (cleanupError: any) {
        this.logger?.error(
          '[NetaClaw] failed to persist assistant placeholder cleanup: %s',
          cleanupError?.message || String(cleanupError)
        );
      }
      throw error;
    }

    const nextMetadata = {
      ...(result.metadata || {}),
      inFlight: false,
      ...(subagentBatches.length > 0
        ? { subagentBatches: this.cloneSubagentBatches(subagentBatches) }
        : {}),
    };
    await this.sessionService.updateMessage(assistantPlaceholder.id, {
      content: result.finalContent,
      thinking: result.thinking,
      metadata: nextMetadata,
    });

    return {
      sessionId,
      finalContent: result.finalContent,
      content: result.finalContent,
      usage: result.usage,
      toolCallCount: result.toolCallCount ?? 0,
      metadata: nextMetadata,
      assistantMessageId: assistantPlaceholder.id,
      subagentBatches: this.cloneSubagentBatches(subagentBatches),
    };
  }

  async executeSessionTreeChat(params: {
    sessionId?: string;
    requestedAgentName: string;
    message: string;
    userMessageMetadata?: Record<string, unknown>;
    agentId?: number;
    userId?: string;
    cwd?: string;
    sessionMetadata?: Record<string, unknown>;
    realtime?: SessionTreeChatRealtimeSink;
    runner: (params: SessionTreeChatOrchestratorRunnerParams) => Promise<ChatOrchestratorRunnerResult>;
  }) {
    const sessionId = await this.sessionService.getOrCreateSession(
      params.sessionId,
      params.requestedAgentName,
      params.agentId,
      params.userId,
      params.cwd,
      params.sessionMetadata,
    );

    const snapshot = await this.sessionService.getRuntimeSnapshot(sessionId, params.agentId);
    const savedUserEntry = await this.sessionService.appendUserMessageEntry(sessionId, {
      parentEntryId: snapshot.session.leafEntryId,
      message: params.message,
      agentId: params.agentId,
      metadata: params.userMessageMetadata,
    });
    params.realtime?.onEntryAdded?.({
      sessionId,
      entry: savedUserEntry.entry,
      session: savedUserEntry.session,
    });
    params.realtime?.onLeafChanged?.({
      sessionId,
      leafEntryId: savedUserEntry.session.leafEntryId,
      session: savedUserEntry.session,
    });
    const assistantPlaceholder = await this.sessionService.appendAssistantPlaceholderEntry(sessionId, {
      parentEntryId: savedUserEntry.id,
      agentId: params.agentId,
    });
    params.realtime?.onEntryAdded?.({
      sessionId,
      entry: assistantPlaceholder.entry,
      session: assistantPlaceholder.session,
    });
    params.realtime?.onLeafChanged?.({
      sessionId,
      leafEntryId: assistantPlaceholder.session.leafEntryId,
      session: assistantPlaceholder.session,
    });
    params.realtime?.onAssistantStreamStart?.({
      sessionId,
      assistantEntryId: assistantPlaceholder.id,
    });

    const subagentBatches: ChatOrchestratorSubagentBatch[] = [];
    const subagents = this.createSubagentContext(sessionId, 0, subagentBatches, params.realtime);
    const sessionTreeSubagentKey = `${sessionId}:${assistantPlaceholder.id}`;
    this.sessionTreeSubagentParentEntryIds.set(sessionTreeSubagentKey, assistantPlaceholder.id);

    let result: ChatOrchestratorRunnerResult;
    try {
      result = await params.runner({
        sessionId,
        userEntryId: savedUserEntry.id,
        assistantEntryId: assistantPlaceholder.id,
        userMessage: params.message,
        history: snapshot.runtimeContext.messages.filter(message =>
          message.role === 'system' || message.role === 'user' || message.role === 'assistant' || message.role === 'tool',
        ) as LLMMessage[],
        subagents,
        onAssistantDelta: text => {
          params.realtime?.onAssistantStreamDelta?.({
            sessionId,
            assistantEntryId: assistantPlaceholder.id,
            text,
          });
        },
      });
    } catch (error: any) {
      this.sessionTreeSubagentParentEntryIds.delete(sessionTreeSubagentKey);
      const errorMessage = error?.message || 'Agent execution failed';
      const finalized = await this.sessionService.finalizeAssistantEntry(sessionId, assistantPlaceholder.id, {
        content: errorMessage,
        metadata: {
          ...(error?.runtimeMetadata || {}),
          inFlight: false,
          error: {
            message: errorMessage,
            name: error?.name || 'Error',
          },
        },
        agentId: params.agentId,
      });
      params.realtime?.onEntryUpdated?.({ sessionId, entry: finalized.entry, session: finalized.session });
      params.realtime?.onAssistantStreamEnd?.({ sessionId, assistantEntryId: assistantPlaceholder.id });
      throw error;
    }

    const finalized = await this.sessionService.finalizeAssistantEntry(sessionId, assistantPlaceholder.id, {
      content: result.finalContent,
      thinking: result.thinking,
      metadata: {
        ...(result.metadata || {}),
        inFlight: false,
      },
      agentId: params.agentId,
    });
    this.sessionTreeSubagentParentEntryIds.delete(sessionTreeSubagentKey);
    params.realtime?.onEntryUpdated?.({ sessionId, entry: finalized.entry, session: finalized.session });
    params.realtime?.onAssistantStreamEnd?.({ sessionId, assistantEntryId: assistantPlaceholder.id });

    return {
      sessionId,
      finalContent: result.finalContent,
      content: result.finalContent,
      usage: result.usage,
      toolCallCount: result.toolCallCount ?? 0,
      metadata: {
        ...(result.metadata || {}),
        inFlight: false,
      },
      assistantEntryId: assistantPlaceholder.id,
      subagentBatches: this.cloneSubagentBatches(subagentBatches),
    };
  }

  private excludeMessageById(history: LLMMessageWithId[], excludedId: number): LLMMessage[] {
    return history
      .filter(message => message.id !== excludedId)
      .map(({ role, content, toolCalls, toolCallId }) => ({
        role,
        content,
        toolCalls,
        toolCallId,
      }));
  }

  private createSubagentContext(
    sessionId: string,
    parentMessageId: number,
    batches: ChatOrchestratorSubagentBatch[],
    realtimeSink?: SessionTreeChatRealtimeSink,
  ): ChatOrchestratorSubagentContext {
    const upsertBatch = (
      batchId: string,
      patch?: Partial<ChatOrchestratorSubagentBatch>
    ): ChatOrchestratorSubagentBatch => {
      let batch = batches.find(item => item.batchId === batchId);
      if (!batch) {
        batch = {
          batchId,
          sessionId,
          parentMessageId,
          mode: 'single',
          tasks: [],
          status: 'running',
          results: [],
        };
        batches.push(batch);
      }

      if (patch) {
        Object.assign(batch, patch);
      }

      return batch;
    };

    return {
      sessionId,
      parentMessageId,
      batches,
      onBatchStart: event => {
        const batch = upsertBatch(event.batchId, {
          sessionId: event.sessionId,
          parentMessageId: event.parentMessageId,
          mode: event.mode,
          tasks: event.tasks ? event.tasks.map(task => ({ ...task })) : [],
          status: 'running',
        });
        if (event.toolRuntimeRoutes?.length) {
          batch.metadata = this.buildSubagentBatchMetadata(batch, undefined, {
            toolRuntimeRoutes: event.toolRuntimeRoutes.map(route => ({ ...route })),
          });
        }
        const treeParentEntryId = this.sessionTreeSubagentParentEntryIds.get(`${event.sessionId}:${event.parentMessageId}`);
        if (treeParentEntryId) {
          void this.sessionService.appendSubagentBatchEntry(event.sessionId, {
            parentEntryId: treeParentEntryId,
            batchId: event.batchId,
            mode: event.mode,
            tasks: (event.tasks ?? []).map(task => ({ ...task })),
            status: 'running',
            agentId: undefined,
          }).then(added => {
            this.sessionTreeSubagentBatchEntryIds.set(`${event.sessionId}:${event.batchId}`, added.id);
            realtimeSink?.onEntryAdded?.({
              sessionId: event.sessionId,
              entry: added.entry,
              session: added.session,
            });
            if (batch.metadata) {
              void this.flushSubagentBatchEntryUpdate({
                sessionId: event.sessionId,
                batch,
                batchEntryId: added.id,
                realtimeSink,
              }).catch(error => {
                this.logger?.error('[NetaClaw] failed to flush subagent batch entry update: %s', error?.message || String(error));
              });
            }
          }).catch(error => {
            this.logger?.error('[NetaClaw] failed to append subagent batch entry: %s', error?.message || String(error));
          });
        }
        realtimeSink?.onSubagentRunStarted?.({
          sessionId: event.sessionId,
          batchId: event.batchId,
          parentEntryId: treeParentEntryId,
          mode: event.mode,
          tasks: event.tasks,
        });
      },
      onUpdate: event => {
        const batch = upsertBatch(event.batchId, {
          sessionId: event.sessionId,
          parentMessageId: event.parentMessageId,
        });
        const record = { ...event.event.subagent };
        const existingIndex = batch.results.findIndex(item =>
          (typeof record.id === 'number' && item.id === record.id)
          || (
            typeof record.sortOrder === 'number'
            && typeof item.sortOrder === 'number'
            && item.sortOrder === record.sortOrder
          ));

        if (existingIndex >= 0) {
          batch.results[existingIndex] = {
            ...batch.results[existingIndex],
            ...record,
          };
        } else {
          batch.results.push(record);
        }
        const treeParentEntryId = this.sessionTreeSubagentParentEntryIds.get(`${event.sessionId}:${event.parentMessageId}`);
        const batchEntryId = this.sessionTreeSubagentBatchEntryIds.get(`${event.sessionId}:${event.batchId}`);
        const nextEvent = {
          sequence: event.sequence,
          timestamp: event.timestamp,
          event: {
            kind: event.event.kind,
            status: event.event.status,
            subagent: record,
          },
        };
        batch.metadata = this.buildSubagentBatchMetadata(batch, nextEvent);
        if (batchEntryId) {
          void this.flushSubagentBatchEntryUpdate({
            sessionId: event.sessionId,
            batch,
            batchEntryId,
            realtimeSink,
          }).catch(error => {
            this.logger?.error('[NetaClaw] failed to update subagent batch entry: %s', error?.message || String(error));
          });
        }
        realtimeSink?.onSubagentEvent?.({
          sessionId: event.sessionId,
          batchId: event.batchId,
          parentEntryId: treeParentEntryId,
          sequence: event.sequence,
          timestamp: event.timestamp,
          event: {
            kind: event.event.kind,
            status: event.event.status,
            subagent: record,
          },
        });
      },
      onDone: event => {
        const batch = upsertBatch(event.batchId, {
          sessionId: event.sessionId,
          parentMessageId: event.parentMessageId,
          status: 'completed',
        });
        batch.results = event.results.map(result => ({ ...result }));
        const treeParentEntryId = this.sessionTreeSubagentParentEntryIds.get(`${event.sessionId}:${event.parentMessageId}`);
        const batchEntryId = this.sessionTreeSubagentBatchEntryIds.get(`${event.sessionId}:${event.batchId}`);
        const completionEvent = {
          sequence: Number.MAX_SAFE_INTEGER,
          timestamp: new Date().toISOString(),
          event: {
            kind: 'status' as const,
            status: 'completed' as const,
            subagent: {
              status: 'completed',
              summary: `Completed ${event.results.length} subagent result(s)`,
            },
          },
        };
        batch.metadata = this.buildSubagentBatchMetadata(batch, completionEvent, {
          finalResults: event.results.map(result => ({ ...result })),
        });
        if (batchEntryId) {
          void this.flushSubagentBatchEntryUpdate({
            sessionId: event.sessionId,
            batch,
            batchEntryId,
            realtimeSink,
          }).catch(error => {
            this.logger?.error('[NetaClaw] failed to finalize subagent batch entry: %s', error?.message || String(error));
          }).finally(() => {
            this.sessionTreeSubagentBatchEntryIds.delete(`${event.sessionId}:${event.batchId}`);
          });
        }
        if (treeParentEntryId) {
          const processEvents: SubagentProcessEventProjection[] = event.results.flatMap(result => {
            const payload = result.resultPayload;
            if (!payload || typeof payload !== 'object' || !Array.isArray((payload as Record<string, unknown>).processEvents)) {
              return [];
            }
            return ((payload as Record<string, unknown>).processEvents as unknown[])
              .map(item => toSubagentProcessEventProjection(item))
              .filter((item): item is SubagentProcessEventProjection => Boolean(item));
          });
          const evidenceSummaries: SubagentEvidenceSummaryProjection[] = event.results
            .map(result => result.resultPayload)
            .filter((payload): payload is Record<string, unknown> => Boolean(payload) && typeof payload === 'object')
            .map(payload => toSubagentEvidenceSummaryProjection(payload.evidenceSummary))
            .filter((summary): summary is SubagentEvidenceSummaryProjection => Boolean(summary));
          void this.sessionService.appendSubagentResultEntry(event.sessionId, {
            parentEntryId: treeParentEntryId,
            batchId: event.batchId,
            results: event.results.map(result => ({ ...result })),
            status: 'completed',
            metadata: (processEvents.length > 0 || evidenceSummaries.length > 0
              ? {
                  finalResults: event.results.map(result => ({ ...result })),
                  ...(processEvents.length > 0 ? { processEvents } : {}),
                  ...(evidenceSummaries.length > 0 ? { evidenceSummaries } : {}),
                }
              : {
                  finalResults: event.results.map(result => ({ ...result })),
                }) as SessionTreeSubagentResultMetadata,
          }).then(added => {
            realtimeSink?.onEntryAdded?.({
              sessionId: event.sessionId,
              entry: added.entry,
              session: added.session,
            });
          }).catch(error => {
            this.logger?.error('[NetaClaw] failed to append subagent result entry: %s', error?.message || String(error));
          });
        }
        realtimeSink?.onSubagentRunCompleted?.({
          sessionId: event.sessionId,
          batchId: event.batchId,
          parentEntryId: treeParentEntryId,
          results: event.results.map(result => ({ ...result })),
        });
      },
    };
  }

  private cloneSubagentBatches(batches: ChatOrchestratorSubagentBatch[]): ChatOrchestratorSubagentBatch[] {
    return batches.map(batch => ({
      ...batch,
      tasks: batch.tasks.map(task => ({ ...task })),
      results: batch.results.map(result => ({ ...result })),
      metadata: batch.metadata ? JSON.parse(JSON.stringify(batch.metadata)) as Record<string, unknown> : undefined,
    }));
  }

}

function toSubagentProcessEventProjection(value: unknown): SubagentProcessEventProjection | null {
  if (!value || typeof value !== 'object' || Array.isArray(value)) {
    return null;
  }
  const record = value as Record<string, unknown>;
  if (typeof record.type !== 'string' || !record.type.trim()) {
    return null;
  }
  return { ...record, type: record.type } as SubagentProcessEventProjection;
}

function toSubagentEvidenceSummaryProjection(value: unknown): SubagentEvidenceSummaryProjection | null {
  if (!value || typeof value !== 'object' || Array.isArray(value)) {
    return null;
  }
  const record = value as Record<string, unknown>;
  if (typeof record.kind !== 'string' || typeof record.title !== 'string' || typeof record.summary !== 'string') {
    return null;
  }
  return {
    ...record,
    kind: record.kind,
    title: record.title,
    summary: record.summary,
  } as SubagentEvidenceSummaryProjection;
}
