import { ulid } from "ulid"
import { Effect, Context, Layer, Cause } from "effect"
import { MessageID, PartID, SessionID } from "./schema"
import { MessageV2 } from "../message/message-v2"
import * as Session from "./session"
import { Agent } from "../../agent/agent"
import { Provider } from "@/provider/provider"
import { ProviderID, ModelID } from "../../provider/schema"
import { Bus } from "../../bus"
import { Permission } from "@/permission"
import { Plugin } from "../../plugin"
import { ToolRegistry } from "@/tool/registry"
import { TaskTool, type TaskPromptOps } from "@/tool/task"
import { EffectBridge } from "@/effect/bridge"
import { NamedError } from "@saeeol/core/util/error"
import { InstanceState } from "@/effect/instance-state"
import { SaeeolSessionPrompt } from "@/saeeol/session/prompt" 
import { SaeeolCostPropagation } from "@/saeeol/session/cost-propagation" 
import * as Log from "@saeeol/core/util/log"

const log = Log.create({ service: "session.subtask" })

export interface SubtaskInput {
  task: MessageV2.SubtaskPart
  model: Provider.Model
  lastUser: MessageV2.User
  sessionID: SessionID
  session: Session.Info
  msgs: MessageV2.WithParts[]
  getModel: (p: ProviderID, m: ModelID, s: SessionID) => Effect.Effect<Provider.Model>
  promptOps: TaskPromptOps
}

export interface Interface {
  readonly handle: (input: SubtaskInput) => Effect.Effect<void>
}

export class Service extends Context.Service<Service, Interface>()("@saeeol/SessionSubtask") {}

export const layer = Layer.effect(
  Service,
  Effect.gen(function* () {
    const sessions = yield* Session.Service
    const agents = yield* Agent.Service
    const bus = yield* Bus.Service
    const plugin = yield* Plugin.Service
    const permission = yield* Permission.Service
    const registry = yield* ToolRegistry.Service
    const bridge = yield* EffectBridge.make()

    const handle = Effect.fn("SessionSubtask.handle")(function* (input: SubtaskInput) {
      const { task, model, lastUser, sessionID, session, msgs } = input
      const ctx = yield* InstanceState.context
      const { task: taskTool } = yield* registry.named()
      const taskModel = task.model ? yield* input.getModel(task.model.providerID, task.model.modelID, sessionID) : model
      const assistantMessage: MessageV2.Assistant = yield* sessions.updateMessage({
        id: MessageID.ascending(),
        role: "assistant",
        parentID: lastUser.id,
        sessionID,
        mode: task.agent,
        agent: task.agent,
        variant: lastUser.model.variant,
        path: { cwd: ctx.directory, root: ctx.worktree },
        cost: 0,
        tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
        modelID: taskModel.id,
        providerID: taskModel.providerID,
        time: { created: Date.now() },
      })
      let part: MessageV2.ToolPart = yield* sessions.updatePart({
        id: PartID.ascending(),
        messageID: assistantMessage.id,
        sessionID: assistantMessage.sessionID,
        type: "tool",
        callID: ulid(),
        tool: TaskTool.id,
        state: {
          status: "running",
          input: {
            prompt: task.prompt,
            description: task.description,
            subagent_type: task.agent,
            command: task.command,
          },
          time: { start: Date.now() },
        },
      })
      const taskArgs = {
        prompt: task.prompt,
        description: task.description,
        subagent_type: task.agent,
        command: task.command,
      }
      yield* plugin.trigger(
        "tool.execute.before",
        { tool: TaskTool.id, sessionID, callID: part.id },
        { args: taskArgs },
      )

      const taskAgent = yield* agents.get(task.agent)
      if (!taskAgent) {
        const available = (yield* agents.list()).filter((a: any) => !a.hidden).map((a: any) => a.name)
        const hint = available.length ? ` Available agents: ${available.join(", ")}` : ""
        const error = new NamedError.Unknown({ message: `Agent not found: "${task.agent}".${hint}` })
        yield* bus.publish(Session.Event.Error, { sessionID, error: error.toObject() })
        throw error
      }

      let error: Error | undefined
      const taskAbort = new AbortController()
      const childID = () => {
        const meta = part.state.status !== "pending" ? part.state.metadata : undefined
        return (meta as { sessionId?: string } | undefined)?.sessionId
      }

      const result = yield* taskTool
        .execute(taskArgs, {
          agent: task.agent,
          messageID: assistantMessage.id,
          sessionID,
          abort: taskAbort.signal,
          callID: part.callID,
          extra: { bypassAgentCheck: true, promptOps: input.promptOps },
          messages: msgs,
          metadata: (val: any) =>
            Effect.gen(function* () {
              part = yield* sessions.updatePart({
                ...part,
                type: "tool",
                state: { ...part.state, ...val },
              } satisfies MessageV2.ToolPart)
            }),
          ask: (req: any) =>
            permission
              .ask({
                ...req,
                sessionID,
                ruleset: Permission.merge(
                  taskAgent.permission,
                  SaeeolSessionPrompt.guardPermissions({ agent: taskAgent, session }),
                ),
                hardRuleset: SaeeolSessionPrompt.hardPermissions({ agent: taskAgent }),
              })
              .pipe(Effect.orDie),
        })
        .pipe(
          Effect.catchCause((cause) => {
            const defect = Cause.squash(cause)
            error = defect instanceof Error ? defect : new Error(String(defect))
            log.error("subtask execution failed", { error, agent: task.agent, description: task.description })
            return Effect.void
          }),
          Effect.onInterrupt(() =>
            Effect.gen(function* () {
              taskAbort.abort()
              assistantMessage.finish = "tool-calls"
              assistantMessage.time.completed = Date.now()
              const cid = childID()
              if (cid) {
                assistantMessage.cost = yield* SaeeolCostPropagation.childCost(sessions, SessionID.make(cid))
              }
              yield* sessions.updateMessage(assistantMessage)
              if (part.state.status === "running") {
                yield* sessions.updatePart({
                  ...part,
                  state: {
                    status: "error",
                    error: "Cancelled",
                    time: { start: part.state.time.start, end: Date.now() },
                    metadata: part.state.metadata,
                    input: part.state.input,
                  },
                } satisfies MessageV2.ToolPart)
              }
            }),
          ),
        )

      const attachments = result?.attachments?.map((attachment: any) => ({
        ...attachment,
        id: PartID.ascending(),
        sessionID,
        messageID: assistantMessage.id,
      }))

      yield* plugin.trigger(
        "tool.execute.after",
        { tool: TaskTool.id, sessionID, callID: part.id, args: taskArgs },
        result,
      )

      assistantMessage.finish = "tool-calls"
      assistantMessage.time.completed = Date.now()
      const cid = result?.metadata?.sessionId ?? childID()
      if (cid) {
        assistantMessage.cost = yield* SaeeolCostPropagation.childCost(sessions, SessionID.make(cid))
      }
      yield* sessions.updateMessage(assistantMessage)

      if (result && part.state.status === "running") {
        yield* sessions.updatePart({
          ...part,
          state: {
            status: "completed",
            input: part.state.input,
            title: result.title,
            metadata: result.metadata,
            output: result.output,
            attachments,
            time: { ...part.state.time, end: Date.now() },
          },
        } satisfies MessageV2.ToolPart)
      }

      if (!result) {
        yield* sessions.updatePart({
          ...part,
          state: {
            status: "error",
            error: error ? `Tool execution failed: ${error.message}` : "Tool execution failed",
            time: {
              start: part.state.status === "running" ? part.state.time.start : Date.now(),
              end: Date.now(),
            },
            metadata: part.state.status === "pending" ? undefined : part.state.metadata,
            input: part.state.input,
          },
        } satisfies MessageV2.ToolPart)
      }

      if (!task.command) return

      const summaryUserMsg: MessageV2.User = {
        id: MessageID.ascending(),
        sessionID,
        role: "user",
        time: { created: Date.now() },
        agent: lastUser.agent,
        model: lastUser.model,
        editorContext: lastUser.editorContext,
      }
      yield* sessions.updateMessage(summaryUserMsg)
      yield* sessions.updatePart({
        id: PartID.ascending(),
        messageID: summaryUserMsg.id,
        sessionID,
        type: "text",
        text: "Summarize the task tool output above and continue with your task.",
        synthetic: true,
      } satisfies MessageV2.TextPart)
    })

    return { handle }
  }),
)

export const defaultLayer = Layer.suspend(() =>
  layer.pipe(
    Layer.provide(Session.defaultLayer),
    Layer.provide(Agent.defaultLayer),
    Layer.provide(Plugin.defaultLayer),
    Layer.provide(Permission.defaultLayer),
    Layer.provide(ToolRegistry.defaultLayer),
    Layer.provide(Provider.defaultLayer),
    Layer.provide(Bus.layer),
  ),
)
