#!/usr/bin/env node
/**
 * Diligence MCP Server
 *
 * Enforces diligent code changes through a worker-reviewer loop.
 * Prevents quick fixes that miss patterns, events, and edge cases.
 *
 * Workflow:
 * 1. User describes task → conversation phase
 * 2. Start work → research phase begins
 * 3. Worker proposes → Reviewer challenges → repeat until approved
 * 4. Approved → implementation phase
 * 5. Complete → back to conversation
 *
 * Project-specific context is loaded from .claude/CODEBASE_CONTEXT.md
 * Generic prompts are bundled with this server.
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { ListToolsRequestSchema, CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { existsSync, readFileSync, writeFileSync, mkdirSync, readdirSync } from 'fs';
import { join, dirname, basename } from 'path';
import { fileURLToPath } from 'url';
import { execSync } from 'child_process';

// =============================================================================
// Configuration
// =============================================================================

const __dirname = dirname(fileURLToPath(import.meta.url));
const MAX_ROUNDS = 5;

/**
 * Find the project root by looking for .git directory.
 * Falls back to current working directory.
 */
function findProjectRoot() {
  try {
    const gitRoot = execSync('git rev-parse --show-toplevel', {
      encoding: 'utf-8',
      stdio: ['pipe', 'pipe', 'pipe']
    }).trim();
    return gitRoot;
  } catch {
    return process.cwd();
  }
}

const PROJECT_ROOT = findProjectRoot();
const STATE_DIR = join(PROJECT_ROOT, '.claude');
const STATE_FILE = join(STATE_DIR, '.diligence-state.json');
const CODEBASE_CONTEXT_FILE = join(STATE_DIR, 'CODEBASE_CONTEXT.md');

// Bundled prompts (shipped with this package)
const WORKER_PROMPT_FILE = join(__dirname, 'WORKER_PROMPT.md');
const REVIEWER_PROMPT_FILE = join(__dirname, 'REVIEWER_PROMPT.md');

// =============================================================================
// State Management
// =============================================================================

/**
 * @typedef {Object} WorkflowState
 * @property {'conversation' | 'researching' | 'approved' | 'implementing'} phase
 * @property {string | null} task
 * @property {number} round
 * @property {string | null} proposal
 * @property {Array<{round: number, feedback: string}>} feedback
 * @property {Array<string>} userContext
 * @property {string | null} approvalReason
 * @property {string | null} startedAt
 */

/** @type {WorkflowState} */
let state = {
  phase: 'conversation',
  task: null,
  round: 0,
  proposal: null,
  feedback: [],
  userContext: [],
  approvalReason: null,
  startedAt: null,
};

function loadState() {
  try {
    if (existsSync(STATE_FILE)) {
      const saved = JSON.parse(readFileSync(STATE_FILE, 'utf-8'));
      state = { ...state, ...saved };
    }
  } catch {
    // Use default state
  }
}

function saveState() {
  try {
    if (!existsSync(STATE_DIR)) {
      mkdirSync(STATE_DIR, { recursive: true });
    }
    writeFileSync(STATE_FILE, JSON.stringify(state, null, 2));
  } catch {
    // Ignore write errors
  }
}

function resetState() {
  state = {
    phase: 'conversation',
    task: null,
    round: 0,
    proposal: null,
    feedback: [],
    userContext: [],
    approvalReason: null,
    startedAt: null,
  };
  saveState();
}

// =============================================================================
// Context Loading
// =============================================================================

// Project-specific context files
const PROJECT_WORKER_CONTEXT = join(STATE_DIR, 'WORKER_CONTEXT.md');
const PROJECT_REVIEWER_CONTEXT = join(STATE_DIR, 'REVIEWER_CONTEXT.md');
const PROJECT_CONTEXT_DIR = join(STATE_DIR, 'context');

function loadFile(filePath) {
  try {
    if (existsSync(filePath)) {
      return readFileSync(filePath, 'utf-8');
    }
  } catch {
    // Ignore
  }
  return '';
}

/**
 * Load all .md files from a directory.
 * @param {string} dirPath
 * @returns {Array<{name: string, content: string}>}
 */
function loadContextDir(dirPath) {
  const files = [];
  try {
    if (existsSync(dirPath)) {
      const entries = readdirSync(dirPath);
      for (const entry of entries) {
        if (entry.endsWith('.md')) {
          const content = loadFile(join(dirPath, entry));
          if (content) {
            files.push({ name: basename(entry, '.md'), content });
          }
        }
      }
    }
  } catch {
    // Ignore
  }
  return files;
}

function getCodebaseContext() {
  let context = '';

  // Main codebase context
  const mainContent = loadFile(CODEBASE_CONTEXT_FILE);
  if (mainContent) {
    context += mainContent;
  } else {
    context += `(No CODEBASE_CONTEXT.md found at ${CODEBASE_CONTEXT_FILE})\n\nCreate this file to provide project-specific architecture and patterns.`;
  }

  // Additional context files from .claude/context/
  const extraFiles = loadContextDir(PROJECT_CONTEXT_DIR);
  if (extraFiles.length > 0) {
    context += '\n\n---\n\n# Additional Context\n\n';
    for (const file of extraFiles) {
      context += `## ${file.name}\n\n${file.content}\n\n`;
    }
  }

  return context;
}

function getWorkerPrompt() {
  // Bundled generic prompt
  let prompt = loadFile(WORKER_PROMPT_FILE);

  // Project-specific additions
  const projectWorker = loadFile(PROJECT_WORKER_CONTEXT);
  if (projectWorker) {
    prompt += '\n\n---\n\n# Project-Specific Worker Guidance\n\n';
    prompt += projectWorker;
  }

  return prompt;
}

function getReviewerPrompt() {
  // Bundled generic prompt
  let prompt = loadFile(REVIEWER_PROMPT_FILE);

  // Project-specific additions
  const projectReviewer = loadFile(PROJECT_REVIEWER_CONTEXT);
  if (projectReviewer) {
    prompt += '\n\n---\n\n# Project-Specific Reviewer Guidance\n\n';
    prompt += projectReviewer;
  }

  return prompt;
}

// =============================================================================
// MCP Server
// =============================================================================

const server = new Server({
  name: 'diligence',
  version: '0.1.0',
}, {
  capabilities: { tools: {} }
});

// =============================================================================
// Tool Definitions
// =============================================================================

server.setRequestHandler(ListToolsRequestSchema, async () => ({
  tools: [
    {
      name: 'status',
      description: 'Get current diligence workflow state (phase, round, task). Check this before taking actions.',
      inputSchema: { type: 'object', properties: {} }
    },
    {
      name: 'start',
      description: 'Start the diligence workflow for a task. Worker will research and propose, Reviewer will verify.',
      inputSchema: {
        type: 'object',
        properties: {
          task: { type: 'string', description: 'Description of the task/bug to fix' }
        },
        required: ['task']
      }
    },
    {
      name: 'propose',
      description: 'Worker submits a proposal with code citations, data flow analysis, and implementation plan.',
      inputSchema: {
        type: 'object',
        properties: {
          proposal: { type: 'string', description: 'Full proposal in markdown with file:line citations' }
        },
        required: ['proposal']
      }
    },
    {
      name: 'review',
      description: 'Reviewer submits APPROVED or NEEDS_WORK after verifying claims by searching the codebase.',
      inputSchema: {
        type: 'object',
        properties: {
          decision: { type: 'string', enum: ['APPROVED', 'NEEDS_WORK'], description: 'Review decision' },
          reasoning: { type: 'string', description: 'Verification results and feedback' }
        },
        required: ['decision', 'reasoning']
      }
    },
    {
      name: 'add_context',
      description: 'Add user feedback to incorporate in the next round.',
      inputSchema: {
        type: 'object',
        properties: {
          context: { type: 'string', description: 'Additional context from user' }
        },
        required: ['context']
      }
    },
    {
      name: 'get_worker_brief',
      description: 'Get the full brief for the Worker: task, codebase context, previous feedback, instructions.',
      inputSchema: { type: 'object', properties: {} }
    },
    {
      name: 'get_reviewer_brief',
      description: 'Get the full brief for the Reviewer: proposal to verify, codebase context, instructions.',
      inputSchema: { type: 'object', properties: {} }
    },
    {
      name: 'approve',
      description: 'User manually approves, bypassing further review. Use when proposal is clearly correct.',
      inputSchema: {
        type: 'object',
        properties: {
          reason: { type: 'string', description: 'Why approving manually' }
        },
        required: ['reason']
      }
    },
    {
      name: 'implement',
      description: 'Begin implementation phase. Only allowed after approval.',
      inputSchema: { type: 'object', properties: {} }
    },
    {
      name: 'complete',
      description: 'Mark implementation complete. Returns to conversation phase.',
      inputSchema: {
        type: 'object',
        properties: {
          summary: { type: 'string', description: 'Summary of changes made' }
        },
        required: ['summary']
      }
    },
    {
      name: 'abort',
      description: 'Cancel workflow and reset to conversation.',
      inputSchema: {
        type: 'object',
        properties: {
          reason: { type: 'string', description: 'Why aborting' }
        }
      }
    }
  ]
}));

// =============================================================================
// Tool Handlers
// =============================================================================

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  try {
    switch (name) {
      case 'status': {
        const lines = [
          `Phase: ${state.phase}`,
          `Task: ${state.task || '(none)'}`,
          `Round: ${state.round}/${MAX_ROUNDS}`,
          `Proposal: ${state.proposal ? `${state.proposal.length} chars` : '(none)'}`,
          `Feedback rounds: ${state.feedback.length}`,
          state.userContext.length ? `User context: ${state.userContext.length} items` : null,
          state.approvalReason ? `Approved: ${state.approvalReason.slice(0, 80)}...` : null,
          `Project: ${PROJECT_ROOT}`,
        ].filter(Boolean);
        return { content: [{ type: 'text', text: lines.join('\n') }] };
      }

      case 'start': {
        if (state.phase !== 'conversation') {
          return { content: [{ type: 'text', text: `Already in "${state.phase}" phase. Use abort first.` }], isError: true };
        }
        state.phase = 'researching';
        state.task = args.task;
        state.round = 1;
        state.proposal = null;
        state.feedback = [];
        state.userContext = [];
        state.approvalReason = null;
        state.startedAt = new Date().toISOString();
        saveState();
        return { content: [{ type: 'text', text: `Started: "${args.task}"\n\nPhase: researching\nRound: 1/${MAX_ROUNDS}\n\nWorker should use get_worker_brief for context, then research and propose.` }] };
      }

      case 'propose': {
        if (state.phase !== 'researching') {
          return { content: [{ type: 'text', text: `Not in researching phase (current: ${state.phase})` }], isError: true };
        }
        state.proposal = args.proposal;
        saveState();
        return { content: [{ type: 'text', text: `Proposal submitted (${args.proposal.length} chars).\n\nRound: ${state.round}/${MAX_ROUNDS}\n\nReviewer should use get_reviewer_brief, verify claims by searching codebase, then review.` }] };
      }

      case 'review': {
        if (state.phase !== 'researching') {
          return { content: [{ type: 'text', text: `Not in researching phase (current: ${state.phase})` }], isError: true };
        }
        if (!state.proposal) {
          return { content: [{ type: 'text', text: 'No proposal to review yet' }], isError: true };
        }

        const { decision, reasoning } = args;

        if (decision === 'APPROVED') {
          state.phase = 'approved';
          state.approvalReason = reasoning;
          saveState();
          return { content: [{ type: 'text', text: `APPROVED after ${state.round} round(s).\n\n${reasoning.slice(0, 500)}\n\nUse implement to begin making changes.` }] };
        }

        if (decision === 'NEEDS_WORK') {
          state.feedback.push({ round: state.round, feedback: reasoning });
          state.round++;
          state.userContext = [];

          if (state.round > MAX_ROUNDS) {
            const lastFeedback = reasoning;
            resetState();
            return { content: [{ type: 'text', text: `MAX ROUNDS (${MAX_ROUNDS}) reached. Workflow reset.\n\nLast feedback: ${lastFeedback.slice(0, 500)}\n\nBreak down the task or discuss with user.` }] };
          }

          saveState();
          return { content: [{ type: 'text', text: `NEEDS_WORK - Round ${state.round}/${MAX_ROUNDS}\n\n${reasoning}\n\nWorker should revise addressing this feedback.` }] };
        }

        return { content: [{ type: 'text', text: `Invalid decision: ${decision}` }], isError: true };
      }

      case 'add_context': {
        if (state.phase !== 'researching') {
          return { content: [{ type: 'text', text: `Not in researching phase` }], isError: true };
        }
        state.userContext.push(args.context);
        saveState();
        return { content: [{ type: 'text', text: `Context added. ${state.userContext.length} item(s) pending.` }] };
      }

      case 'get_worker_brief': {
        let brief = '# Worker Brief\n\n';
        brief += `## Task\n\n${state.task || '(No task - use start first)'}\n\n`;
        brief += `## Round\n\n${state.round}/${MAX_ROUNDS}\n\n`;

        if (state.feedback.length > 0) {
          brief += '## Previous Feedback (Address ALL)\n\n';
          for (const f of state.feedback) {
            brief += `### Round ${f.round}\n\n${f.feedback}\n\n`;
          }
        }

        if (state.userContext.length > 0) {
          brief += '## User Context\n\n';
          state.userContext.forEach(c => brief += `- ${c}\n`);
          brief += '\n';
        }

        brief += '## Codebase Context\n\n';
        brief += getCodebaseContext();
        brief += '\n\n';

        const workerPrompt = getWorkerPrompt();
        if (workerPrompt) {
          brief += '## Instructions\n\n';
          brief += workerPrompt;
        }

        return { content: [{ type: 'text', text: brief }] };
      }

      case 'get_reviewer_brief': {
        let brief = '# Reviewer Brief\n\n';
        brief += `## Task\n\n${state.task || '(No task)'}\n\n`;
        brief += `## Round\n\n${state.round}/${MAX_ROUNDS}\n\n`;

        if (state.proposal) {
          brief += '## Proposal to Review\n\n';
          brief += state.proposal;
          brief += '\n\n';
        } else {
          brief += '## Proposal\n\n(No proposal yet)\n\n';
        }

        if (state.feedback.length > 0) {
          brief += '## Previous Feedback\n\n';
          for (const f of state.feedback) {
            brief += `### Round ${f.round}\n\n${f.feedback}\n\n`;
          }
        }

        brief += '## Codebase Context\n\n';
        brief += getCodebaseContext();
        brief += '\n\n';

        const reviewerPrompt = getReviewerPrompt();
        if (reviewerPrompt) {
          brief += '## Instructions\n\n';
          brief += reviewerPrompt;
        }

        return { content: [{ type: 'text', text: brief }] };
      }

      case 'approve': {
        if (state.phase !== 'researching') {
          return { content: [{ type: 'text', text: `Not in researching phase` }], isError: true };
        }
        if (!state.proposal) {
          return { content: [{ type: 'text', text: 'No proposal to approve' }], isError: true };
        }
        state.phase = 'approved';
        state.approvalReason = `USER: ${args.reason}`;
        saveState();
        return { content: [{ type: 'text', text: `Manually approved.\n\nUse implement to begin.` }] };
      }

      case 'implement': {
        if (state.phase !== 'approved') {
          return { content: [{ type: 'text', text: `Not approved yet (current: ${state.phase})` }], isError: true };
        }
        state.phase = 'implementing';
        saveState();
        return { content: [{ type: 'text', text: `Implementation phase.\n\nTask: ${state.task}\n\nMake changes following the approved proposal.` }] };
      }

      case 'complete': {
        if (state.phase !== 'implementing') {
          return { content: [{ type: 'text', text: `Not in implementing phase` }], isError: true };
        }
        const task = state.task;
        const rounds = state.round;
        resetState();
        return { content: [{ type: 'text', text: `Complete!\n\nTask: ${task}\nRounds: ${rounds}\nSummary: ${args.summary}\n\nReset to conversation.` }] };
      }

      case 'abort': {
        const prev = { phase: state.phase, task: state.task };
        resetState();
        return { content: [{ type: 'text', text: `Aborted.\n\nWas: ${prev.phase}, "${prev.task || ''}"\nReason: ${args.reason || '(none)'}\n\nReset to conversation.` }] };
      }

      default:
        throw new Error(`Unknown tool: ${name}`);
    }
  } catch (e) {
    return { content: [{ type: 'text', text: `Error: ${e.message}` }], isError: true };
  }
});

// =============================================================================
// Lifecycle
// =============================================================================

process.on('SIGTERM', () => { saveState(); process.exit(0); });
process.on('SIGINT', () => { saveState(); process.exit(0); });

loadState();
const transport = new StdioServerTransport();
await server.connect(transport);