package sidekick

import (
  "context"
  "encoding/json"
  "fmt"
  "os"
  "strings"
)

type Sidekick struct {
  Config              AgentConfig
  Model               ModelConfig
  ToolMapping         map[string]string
  ToolRegistry        map[string]ToolDefinition
  StreamHandler       func(string)
  ReasoningHandler    func(string)
  IntermediateHandler func(string)
}

func NewSidekick(config AgentConfig, model ModelConfig, extTools []ToolDefinition) *Sidekick {
  model.Resolve()
  s := &Sidekick{
    Config:       config,
    Model:        model,
    ToolMapping:  make(map[string]string),
    ToolRegistry: make(map[string]ToolDefinition),
  }
  for name, tool := range DefaultInternalTools() {
    if name == "shell_exec" && !s.Config.EnableShellExec {
      continue
    }
    s.registerTool(name, tool)
  }
  for _, tool := range extTools {
    if len(s.Config.Toolsets) > 0 && !contains(s.Config.Toolsets, tool.Toolset) {
      continue
    }
    s.registerTool(tool.Name, tool)
  }
  return s
}

func (s *Sidekick) registerTool(realName string, tool ToolDefinition) {
  sanitized := SanitizeToolName(realName)
  s.ToolMapping[sanitized] = realName
  s.ToolRegistry[realName] = tool
}

func (s *Sidekick) constructSystemMessage(agentPool map[string]*Sidekick) string {
  var sb strings.Builder
  sb.WriteString(s.Config.SystemPrompt + "\n\n")
  sb.WriteString(fmt.Sprintf("Role: %s\n\n", s.Config.Role))
  if len(s.Config.Goals) > 0 {
    sb.WriteString("Goals:\n")
    for _, g := range s.Config.Goals {
      sb.WriteString(fmt.Sprintf("- %s\n", g))
    }
    sb.WriteString("\n")
  }
  if len(s.ToolRegistry) > 0 {
    sb.WriteString("Available Tools:\n")
    for _, t := range s.ToolRegistry {
      sb.WriteString(fmt.Sprintf("- %s: %s\n", SanitizeToolName(t.Name), t.Description))
    }
    sb.WriteString("\n")
  }
  if len(s.Config.Subagents) > 0 {
    sb.WriteString("Available Subagents:\n")
    for _, sid := range s.Config.Subagents {
      if sub, exists := agentPool[sid]; exists {
        sb.WriteString(fmt.Sprintf("- %s: %s\n", sid, sub.Config.Description))
      }
    }
    sb.WriteString("\n")
  }
  sb.WriteString("Action Format:\nJSON object with action. Example:\n")
  sb.WriteString(`{"action": "TOOL_CALL", "params": {"tool_name": "x", "arguments": {}}, "reasoning": "..."}` +
    "\n")
  return sb.String()
}

func (s *Sidekick) Run(ctx context.Context, inCtx []Message, pool map[string]*Sidekick) (string, error) {
  history := make([]Message, len(inCtx))
  copy(history, inCtx)
  if len(history) == 0 || history[0].Role != MessageRoleSystem {
    history = append([]Message{{Role: MessageRoleSystem, Content: s.constructSystemMessage(pool)}}, history...)
  }
  var tFiles []string
  defer func() {
    for _, f := range tFiles {
      os.Remove(f)
    }
  }()
  iters := s.Config.MaxIterations
  if iters <= 0 {
    iters = 10
  }
  for i := 0; i < iters; i++ {
    done, res, err := s.runStep(ctx, &history, pool, &tFiles)
    if err != nil {
      return "", err
    }
    if done {
      return res, nil
    }
  }
  return "Max iterations reached.", nil
}

func (s *Sidekick) logIntermediate(msg Message) {
  if !s.Config.LogIntermediate || s.IntermediateHandler == nil {
    return
  }
  if msg.Reasoning != "" {
    s.IntermediateHandler(fmt.Sprintf("Reasoning: %s", msg.Reasoning))
  }
  if msg.Content != "" {
    s.IntermediateHandler(fmt.Sprintf("LLM Output: %s", msg.Content))
  } else if len(msg.ToolCalls) > 0 {
    s.IntermediateHandler(fmt.Sprintf("Tool Call: %s", msg.ToolCalls[0].Function.Name))
  }
}

func (s *Sidekick) runStep(ctx context.Context, history *[]Message,
  pool map[string]*Sidekick, tFiles *[]string) (bool, string, error) {
  var sc, rc func(string)
  if s.Config.Streaming {
    if s.StreamHandler != nil {
      sc = s.StreamHandler
    }
    if s.ReasoningHandler != nil {
      rc = s.ReasoningHandler
    }
  }
  msg, err := CallLLM(ctx, s.Model, *history, s.ToolRegistry, sc, rc)
  if err != nil {
    return false, "", fmt.Errorf("LLM call failed: %w", err)
  }
  s.logIntermediate(msg)
  *history = append(*history, msg)

  act, params, parsed, err := s.parseAction(msg)
  if err != nil {
    if !parsed {
      act, params = ActionTypeRespond, map[string]interface{}{"response": msg.Content}
    } else {
      *history = append(*history, Message{Role: MessageRoleUser, Content: fmt.Sprintf("Error: %v", err)})
      return false, "", nil
    }
  }
  if act == ActionTypeRespond {
    if r, ok := params["response"].(string); ok {
      return true, r, nil
    }
    return true, msg.Content, nil
  }
  res, err := s.dispatch(ctx, act, params, pool, tFiles)
  if err != nil {
    res = fmt.Sprintf("Error: %v", err)
  }
  *history = append(*history, Message{Role: MessageRoleUser, Content: res})
  return false, "", nil
}

func (s *Sidekick) dispatch(ctx context.Context, a ActionType, p map[string]interface{},
  pl map[string]*Sidekick, tf *[]string) (string, error) {
  if a == ActionTypeToolCall {
    tName, _ := p["tool_name"].(string)
    args, _ := p["arguments"].(map[string]interface{})
    if aStr, ok := p["arguments"].(string); ok && args == nil {
      args, _ = ParseArgs(aStr)
    }
    if s.Config.LogIntermediate && s.IntermediateHandler != nil {
      argsJSON, _ := json.Marshal(args)
      s.IntermediateHandler(fmt.Sprintf("Tool Call: %s(%s)", tName, string(argsJSON)))
    }
    res, err := s.executeTool(ctx, tName, args)
    if err != nil {
      return "", err
    }
    if s.Config.LogIntermediate && s.IntermediateHandler != nil {
      truncRes := res
      if len(truncRes) > 200 {
        truncRes = truncRes[:200] + "..."
      }
      s.IntermediateHandler(fmt.Sprintf("Tool Result: %s", truncRes))
    }
    return BufferGate(res, tf, s.Config.ToolResponseThreshold)
  }
  if a == ActionTypeDelegate {
    sid, _ := p["subagent_id"].(string)
    task, _ := p["task"].(string)
    return s.executeDelegation(ctx, sid, task, pl)
  }
  return "", fmt.Errorf("unknown action: %s", a)
}

func (s *Sidekick) parseAction(msg Message) (ActionType, map[string]interface{}, bool, error) {
  if len(msg.ToolCalls) > 0 {
    c := msg.ToolCalls[0]
    return ActionTypeToolCall,
      map[string]interface{}{"tool_name": c.Function.Name, "arguments": c.Function.Arguments}, true, nil
  }
  var env ActionEnvelope
  c := strings.TrimSpace(msg.Content)
  if strings.HasPrefix(c, "```json") && strings.HasSuffix(c, "```") {
    c = strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(c, "```json"), "```"))
  }
  if err := json.Unmarshal([]byte(c), &env); err != nil {
    return "", nil, false, err
  }
  return env.Action, env.Params, true, nil
}

func (s *Sidekick) executeTool(ctx context.Context, sName string, args map[string]interface{}) (string, error) {
  rName, ok := s.ToolMapping[sName]
  if !ok {
    return "", fmt.Errorf("tool %s not found", sName)
  }
  tDef, ok := s.ToolRegistry[rName]
  if !ok {
    return "", fmt.Errorf("tool %s not registered", rName)
  }
  if tDef.Internal {
    if tDef.Handler == nil {
      return "", fmt.Errorf("missing handler")
    }
    return tDef.Handler(args)
  }
  return CallExternalTool(ctx, tDef.Toolset, rName, args)
}

func (s *Sidekick) executeDelegation(ctx context.Context, sid string,
  task string, pool map[string]*Sidekick) (string, error) {
  if !contains(s.Config.Subagents, sid) {
    return "", fmt.Errorf("subagent %s not allowed", sid)
  }
  sub, ok := pool[sid]
  if !ok {
    return "", fmt.Errorf("subagent %s not found", sid)
  }
  if s.Config.LogIntermediate && s.IntermediateHandler != nil {
    s.IntermediateHandler(fmt.Sprintf("Delegating to %s: %s", sid, task))
  }
  res, err := sub.Run(ctx, []Message{{Role: MessageRoleUser, Content: task}}, pool)
  if err == nil && s.Config.LogIntermediate && s.IntermediateHandler != nil {
    truncRes := res
    if len(truncRes) > 200 {
      truncRes = truncRes[:200] + "..."
    }
    s.IntermediateHandler(fmt.Sprintf("Subagent %s Result: %s", sid, truncRes))
  }
  return res, err
}

func contains(s []string, v string) bool {
  for _, i := range s {
    if i == v {
      return true
    }
  }
  return false
}