import * as fs from 'node:fs';
import * as path from 'node:path';
import * as os from 'node:os';
import { getLaceDir } from '../config/lace-dir';
import { asSessionId } from '@lace/ent-protocol';
import {
  deriveNextEventSeqAcrossSessionFiles,
  invalidatePersonaCache,
  repairOrphanTurnStarts,
  summarizeDurableEvents,
} from './event-log';
import { withSessionLock } from './session-lock';
import { reconcileHead } from './seq-head';
import { atomicWriteJson } from './atomic-write';
import { SessionStorageError } from '../errors/agent-errors';
import type { RuntimeExecutionBinding } from '../tools/runtime/types';

/**
 * MCP server entry as stored in state.json. Extends the protocol's
 * McpServerConfig shape with a lace-internal `source` tag identifying who
 * owns the entry. The tag is purely lace-side metadata: it never crosses the
 * embedder wire boundary and is added by lace itself when storing entries
 * supplied through embedder calls (session/new, session/load, session/resume,
 * session/fork) versus user-facing calls (ent/mcp/servers/upsert).
 *
 * Missing `source` is treated as `embedder` for migration compatibility with
 * state.json files written before this field existed.
 */
export type McpServerSource = 'embedder' | 'user' | 'persona';

export type StoredMcpServer = {
  name: string;
  command: string;
  args?: string[];
  env?: Record<string, string>;
  transport?: 'stdio' | 'sse' | 'http';
  enabled?: boolean;
  tools?: Record<string, 'allow' | 'ask' | 'deny' | 'disable'>;
  source?: McpServerSource;
};

export type SessionMeta = {
  sessionId: string;
  workDir: string;
  created: string;
  persona?: string;
  parent?: {
    sessionId: string;
    jobId: string;
    personaName?: string;
  };
};

export type SessionState = {
  nextEventSeq: number;
  nextStreamSeq: number;
  /** Accumulated token usage cost in USD for this session */
  sessionCostUsd?: number;
  /** Token usage tracking */
  tokenUsage?: {
    totalInputTokens: number;
    totalOutputTokens: number;
  };
  /**
   * The highest breakpoint `at` value that has been fired in the current
   * ascending pressure run. Reset to 0 when pressure drops below all breakpoints.
   * Undefined on disk is treated as 0 (new sessions / old state.json files).
   */
  highestFiredBreakpointAt?: number;
  config?: {
    executionMode?: 'plan' | 'execute';
    approvalMode?:
      | 'ask'
      | 'approveReads'
      | 'approveEdits'
      | 'approve'
      | 'deny'
      | 'dangerouslySkipPermissions';
    connectionId?: string;
    modelId?: string;
    maxBudgetUsd?: number;
    maxThinkingTokens?: number;
    environment?: Record<string, string>;
    runtimeBinding?: RuntimeExecutionBinding;
    mcpServers?: StoredMcpServer[];
    /** Allowlist of tool names for this session; undefined means no scope filter. */
    toolScope?: string[];
    /** Persona name used to render the session's system prompt. */
    personaName?: string;
    /** Host-only credential broker socket for this session's role (Part B). */
    credentialBrokerSocket?: string;
  };
};

export type LoadedSession = {
  meta: SessionMeta;
  dir: string;
  state: SessionState;
  /**
   * True when this open performed crash recovery — i.e. the prior process died
   * mid-turn and `repairOrphanTurnStarts` synthesized at least one `turn_end`.
   * The session-open handler uses this to inject a "you crashed and restarted"
   * notice so the agent knows its last turn was cut off. Always false unless
   * `repairOrphanTurnStarts` was requested.
   */
  recoveredFromCrash: boolean;
};

export function agentSessionsDir(): string {
  const override = process.env.LACE_SESSION_DIR?.trim();
  const candidates: string[] = [];

  if (override) {
    candidates.push(override);
  } else {
    candidates.push(path.join(getLaceDir(), 'agent-sessions'));

    const xdg = process.env.XDG_STATE_HOME?.trim();
    if (xdg) candidates.push(path.join(xdg, 'lace', 'agent-sessions'));

    const home = process.env.HOME?.trim();
    if (home) candidates.push(path.join(home, '.local', 'state', 'lace', 'agent-sessions'));

    candidates.push(path.join(os.tmpdir(), 'lace', 'agent-sessions'));
  }

  let lastError: unknown = undefined;
  for (const sessionsDir of candidates) {
    try {
      fs.mkdirSync(sessionsDir, { recursive: true, mode: 0o700 });
      return sessionsDir;
    } catch (error) {
      lastError = error;
    }
  }

  const msg = lastError instanceof Error ? lastError.message : String(lastError);
  throw new SessionStorageError(`Session storage unavailable: ${msg}`, candidates[0]);
}

export function getSessionDir(sessionId: string): string {
  return path.join(agentSessionsDir(), asSessionId(sessionId));
}

export function readSessionMeta(sessionDir: string): SessionMeta {
  const metaPath = path.join(sessionDir, 'meta.json');
  const meta = JSON.parse(fs.readFileSync(metaPath, 'utf8')) as SessionMeta;
  asSessionId(meta.sessionId);
  return meta;
}

export function writeSessionMeta(sessionDir: string, meta: SessionMeta): void {
  fs.mkdirSync(sessionDir, { recursive: true });
  atomicWriteJson(path.join(sessionDir, 'meta.json'), meta, { mode: 0o600 });
  // Drop any stale persona cache entry from event-log so the next
  // appendDurableEvent picks up the canonical persona instead of whatever
  // was visible during an earlier read (commonly null, see C4).
  invalidatePersonaCache(sessionDir);
}

export function readSessionState(sessionDir: string): SessionState {
  const statePath = path.join(sessionDir, 'state.json');
  let raw: string;
  try {
    raw = fs.readFileSync(statePath, 'utf8');
  } catch (err) {
    // ENOENT is the genuine new-session case: no state.json yet → defaults.
    // ANY other read error (EACCES, EISDIR, I/O) means an EXISTING session's
    // state is unreadable. Returning a config-less default here is dangerous:
    // loadSession would derive nextEventSeq from the event log, see a mismatch,
    // and write that default back over the real state — silently dropping
    // modelId/connectionId and wedging every future turn. Fail loud instead.
    if ((err as NodeJS.ErrnoException).code === 'ENOENT') {
      return { nextEventSeq: 1, nextStreamSeq: 1 };
    }
    throw err;
  }
  let parsed: Partial<SessionState>;
  try {
    parsed = JSON.parse(raw) as Partial<SessionState>;
  } catch (err) {
    // The file exists but is not valid JSON. Writes are atomic (temp + rename),
    // so a torn read should not happen in normal operation; a parse failure
    // means real corruption. Do NOT collapse to the new-session default — that
    // would let a subsequent write clobber recoverable state. Surface it.
    throw new SessionStorageError(
      `state.json for session ${path.basename(sessionDir)} is corrupt and cannot be parsed: ${
        err instanceof Error ? err.message : String(err)
      }`,
      statePath
    );
  }
  return {
    nextEventSeq: typeof parsed.nextEventSeq === 'number' ? parsed.nextEventSeq : 1,
    nextStreamSeq: typeof parsed.nextStreamSeq === 'number' ? parsed.nextStreamSeq : 1,
    sessionCostUsd: typeof parsed.sessionCostUsd === 'number' ? parsed.sessionCostUsd : undefined,
    tokenUsage:
      typeof parsed.tokenUsage === 'object' && parsed.tokenUsage
        ? {
            totalInputTokens: parsed.tokenUsage.totalInputTokens ?? 0,
            totalOutputTokens: parsed.tokenUsage.totalOutputTokens ?? 0,
          }
        : undefined,
    config:
      typeof parsed.config === 'object' && parsed.config
        ? (parsed.config as SessionState['config'])
        : undefined,
    highestFiredBreakpointAt:
      typeof parsed.highestFiredBreakpointAt === 'number'
        ? parsed.highestFiredBreakpointAt
        : undefined,
  };
}

export function writeSessionState(sessionDir: string, state: SessionState): void {
  fs.mkdirSync(sessionDir, { recursive: true });
  atomicWriteJson(path.join(sessionDir, 'state.json'), state, { mode: 0o600 });
}

export function ensureSessionFiles(sessionDir: string): void {
  // Post-migration, new sessions write events under transcripts/<persona>/<date>/
  // and never touch the legacy events.jsonl path. We deliberately do not
  // pre-create an empty events.jsonl here so readAllSessionEventLines doesn't
  // open a zero-line file on every read. Pre-migration sessions keep their
  // legacy file; the dual-read path picks it up when it exists.
  fs.mkdirSync(sessionDir, { recursive: true });
}

// ACP-aligned: workDir→cwd, lastActive→updatedAt
export function listSessions(cwd?: string): Array<{
  sessionId: string;
  cwd: string;
  created: string;
  updatedAt: string;
  messageCount: number;
}> {
  const sessionsDir = agentSessionsDir();
  const sessionIds = fs
    .readdirSync(sessionsDir, { withFileTypes: true })
    .filter((d) => d.isDirectory())
    .map((d) => d.name);

  return sessionIds
    .map((sessionId) => {
      try {
        const sessionDir = path.join(sessionsDir, sessionId);
        const meta = readSessionMeta(sessionDir);
        const summary = summarizeDurableEvents(sessionDir);
        return {
          sessionId: meta.sessionId,
          cwd: meta.workDir, // ACP field name
          created: meta.created,
          updatedAt: summary.lastActive ?? meta.created, // ACP field name
          messageCount: summary.messageCount,
        };
      } catch {
        return null;
      }
    })
    .filter((s): s is NonNullable<typeof s> => !!s)
    .filter((s) => (cwd ? s.cwd === cwd : true));
}

/**
 * Options for `loadSession`.
 *
 * `repairOrphanTurnStarts` is opt-in because `loadSession` is called both
 * for actual session-open (cold start from disk) AND for mid-flight
 * in-memory refreshes after a write. The crash-recovery scan must only run
 * on a true open — synthesizing a turn_end for the active in-flight turn
 * would corrupt the durable log.
 */
export type LoadSessionOptions = {
  /**
   * When true, scan events.jsonl for orphan turn_starts left by a prior
   * process and synthesize a `turn_end` event with stopReason='process_died'
   * for each. Use only at cold session-open (e.g. session/load,
   * session/resume) before any new event is appended for this session in
   * the current process. Default: false.
   */
  repairOrphanTurnStarts?: boolean;
};

export function loadSession(sessionId: string, options?: LoadSessionOptions): LoadedSession {
  const sessionDir = getSessionDir(sessionId);
  const metaPath = path.join(sessionDir, 'meta.json');
  if (!fs.existsSync(metaPath)) throw new Error('Session not found');

  const meta = readSessionMeta(sessionDir);
  let state = readSessionState(sessionDir);
  ensureSessionFiles(sessionDir);

  // The per-session head file (<sessionDir>/.seq) is the seq authority. On open
  // reconcile it monotonically against the JSONL: head = MAX(readHead()|0,
  // MAX(JSONL)+1). This can only move the head UP, so a stale or lost head can
  // never hand out a seq <= an existing JSONL seq. We do this under the
  // per-session lock so a concurrent appender cannot land between the MAX(JSONL)
  // scan and the head write. `state.nextEventSeq` is advisory/derived only —
  // it no longer drives seq assignment and we do NOT rewrite state.json merely
  // to "repair" it (which, once gaps exist, would rewrite on every open).
  const canonicalSessionId = path.basename(sessionDir);
  const laceDir = getLaceDir();
  const reconciledHead = withSessionLock(sessionDir, () =>
    reconcileHead(
      sessionDir,
      () => deriveNextEventSeqAcrossSessionFiles(laceDir, canonicalSessionId) - 1
    )
  );
  state.nextEventSeq = reconciledHead;

  let recoveredFromCrash = false;
  if (options?.repairOrphanTurnStarts) {
    // Crash recovery: if the prior process died between turn_start and turn_end
    // (SIGKILL, OOM, container restart), synthesize a turn_end with
    // stopReason='process_died' so the durable log honors the invariant that
    // every turn_start has a matching turn_end. Must run BEFORE any new event
    // is appended by the opener so the new turn's eventSeq sits after the
    // synthesized one.
    const repair = repairOrphanTurnStarts(sessionDir, state);
    if (repair.synthesized > 0) {
      state = repair.nextState;
      writeSessionState(sessionDir, state);
      recoveredFromCrash = true;
    }
  }

  return { meta, dir: sessionDir, state, recoveredFromCrash };
}
