import { stringify as tomlStringify } from 'jsr:@std/toml';
import type { ComposeContext, ComposeObject } from '../routines/composeTypes';
import {
  BACKEND_MODEL_KEYS,
  detectBackend,
  addBackendDependency,
  ensureV1Suffix,
  setEnvVars,
} from '../routines/backendIntegration';

const DUMMY_API_KEY = 'sk-harbor';
const LISTEN_PORT = 4200;
const INTERNAL_PORT = 4201;

interface OpenFangConfig {
  provider: string;
  model: string;
  baseUrl: string;
  apiKeyEnv: string;
  apiKey: string;
  decayRate: number;
  startupTimeout: number;
  localBackend: boolean;
}

function deriveApiKeyEnv(provider: string): string {
  if (provider === 'ollama') return '';
  return `${provider.toUpperCase()}_API_KEY`;
}

function generateConfigToml(config: OpenFangConfig): string {
  const doc: Record<string, unknown> = {};

  if (config.apiKey) {
    doc.api_key = config.apiKey;
  }

  doc.api_listen = `127.0.0.1:${INTERNAL_PORT}`;

  const defaultModel: Record<string, unknown> = {
    provider: config.provider,
    model: config.model,
  };
  if (config.apiKeyEnv) {
    defaultModel.api_key_env = config.apiKeyEnv;
  }
  if (config.baseUrl) {
    defaultModel.base_url = config.baseUrl;
  }
  doc.default_model = defaultModel;

  const memory: Record<string, unknown> = {
    decay_rate: config.decayRate,
  };
  if (config.localBackend && config.provider !== 'ollama') {
    memory.embedding_provider = 'ollama';
  }
  doc.memory = memory;

  return `# Generated by Harbor\n${tomlStringify(doc)}`;
}

/**
 * Generate an inline bash entrypoint that writes config.toml (with deep-merge
 * if one already exists), injects [model] into agent manifests, boots openfang,
 * and waits for the process.
 *
 * All bare $ are escaped as $$ for Docker Compose interpolation.
 */
function generateEntrypointScript(configToml: string, config: OpenFangConfig): string {
  const timeout = config.startupTimeout;

  const modelToml: Record<string, string> = {
    provider: config.provider,
  };
  if (config.model) modelToml.model = config.model;
  if (config.baseUrl) modelToml.base_url = config.baseUrl;

  const modelSection = tomlStringify({ model: modelToml });

  return `
set -e

mkdir -p /root/.openfang

cat > /tmp/harbor_config.toml << 'HARBOR_CONFIG_EOF'
${configToml}
HARBOR_CONFIG_EOF

if [ -f /root/.openfang/config.toml ]; then
  python3 -c "
import tomllib, copy

with open('/root/.openfang/config.toml', 'rb') as f:
    existing = tomllib.load(f)
with open('/tmp/harbor_config.toml', 'rb') as f:
    harbor = tomllib.load(f)

merged = copy.deepcopy(existing)
for k in ('default_model', 'api_listen', 'api_key', 'memory'):
    if k in harbor:
        if isinstance(merged.get(k), dict) and isinstance(harbor[k], dict):
            merged[k] = {**merged[k], **harbor[k]}
        else:
            merged[k] = harbor[k]

lines = []
scalars = {k: v for k, v in merged.items() if not isinstance(v, dict)}
tables = {k: v for k, v in merged.items() if isinstance(v, dict)}
for k, v in scalars.items():
    lines.append(f'{k} = {repr(v) if isinstance(v, str) else v}')
if scalars:
    lines.append('')
for section, vals in tables.items():
    lines.append(f'[{section}]')
    for k, v in vals.items():
        lines.append(f'{k} = {repr(v) if isinstance(v, str) else v}')
    lines.append('')
with open('/root/.openfang/config.toml', 'w') as f:
    f.write(chr(10).join(lines))
"
else
  cp /tmp/harbor_config.toml /root/.openfang/config.toml
fi

cat >> /tmp/harbor_model.toml << 'HARBOR_MODEL_EOF'
${modelSection}
HARBOR_MODEL_EOF

rm -rf /root/.openfang/agents_runtime
cp -a /root/.openfang/agents /root/.openfang/agents_runtime

for agent_dir in /root/.openfang/agents_runtime/*/; do
  [ -f "$$agent_dir/agent.toml" ] || continue
  python3 -c "
import tomllib, sys, os

agent_path = sys.argv[1]
model_path = '/tmp/harbor_model.toml'

with open(agent_path, 'rb') as f:
    agent = tomllib.load(f)
with open(model_path, 'rb') as f:
    model = tomllib.load(f)

agent['model'] = model.get('model', {})

lines = []
scalars = {k: v for k, v in agent.items() if not isinstance(v, dict)}
tables = {k: v for k, v in agent.items() if isinstance(v, dict)}
for k, v in scalars.items():
    if isinstance(v, list):
        items = ', '.join(repr(i) for i in v)
        lines.append(f'{k} = [{items}]')
    else:
        lines.append(f'{k} = {repr(v) if isinstance(v, str) else v}')
if scalars:
    lines.append('')
for section, vals in tables.items():
    lines.append(f'[{section}]')
    for k, v in vals.items():
        if isinstance(v, list):
            items = ', '.join(repr(i) for i in v)
            lines.append(f'{k} = [{items}]')
        else:
            lines.append(f'{k} = {repr(v) if isinstance(v, str) else v}')
    lines.append('')
with open(agent_path, 'w') as f:
    f.write(chr(10).join(lines))
" "$$agent_dir/agent.toml"
done

openfang start &
OPENFANG_PID=$$!

for i in $$(seq 1 ${timeout}); do
  curl -sf http://127.0.0.1:${INTERNAL_PORT}/api/status > /dev/null 2>&1 && break
  sleep 1
done

for agent_dir in /root/.openfang/agents_runtime/*/; do
  [ -f "$$agent_dir/agent.toml" ] || continue
  openfang agent spawn "$$agent_dir/agent.toml" 2>/dev/null || true
done

exec python3 -c "
import socket, threading, signal, sys
L, R = ${LISTEN_PORT}, ${INTERNAL_PORT}
def relay(s, d):
    try:
        while (b := s.recv(65536)): d.sendall(b)
    except: pass
    finally:
        for x in (s, d):
            try: x.close()
            except: pass
def accept(srv):
    while True:
        c, _ = srv.accept()
        u = socket.socket()
        try:
            u.connect(('127.0.0.1', R))
            threading.Thread(target=relay, args=(c, u), daemon=True).start()
            threading.Thread(target=relay, args=(u, c), daemon=True).start()
        except:
            c.close()
srv = socket.socket()
srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
srv.bind(('0.0.0.0', L))
srv.listen(128)
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
accept(srv)
"
`.trim();
}

export default async function apply(ctx: ComposeContext): Promise<ComposeObject> {
  const { compose, env, services, explicitServices } = ctx;

  if (!compose.services?.openfang) {
    return compose;
  }

  const svc = compose.services.openfang;

  const [
    userProvider,
    userModel,
    userBaseUrl,
    userApiKey,
    decayRateStr,
    startupTimeoutStr,
  ] = await Promise.all([
    env.getValue('openfang.model.provider'),
    env.getValue('openfang.model'),
    env.getValue('openfang.base.url'),
    env.getValue('openfang.api.key'),
    env.getValue('openfang.memory.decay.rate'),
    env.getValue('openfang.startup.timeout'),
  ]);

  const decayRate = parseFloat(decayRateStr) || 0.05;
  const startupTimeout = parseInt(startupTimeoutStr, 10) || 30;

  const backend = detectBackend(services, explicitServices);

  let config: OpenFangConfig;

  if (userBaseUrl) {
    config = {
      provider: userProvider || 'openai',
      model: userModel,
      baseUrl: ensureV1Suffix(userBaseUrl),
      apiKeyEnv: deriveApiKeyEnv(userProvider || 'openai'),
      apiKey: userApiKey,
      decayRate,
      startupTimeout,
      localBackend: false,
    };

    if (backend) {
      addBackendDependency(svc, backend.service);
    }
  } else if (backend) {
    const isOllama = backend.service === 'ollama';
    const provider = isOllama ? 'ollama' : 'openai';

    let model = userModel;
    const modelKey = BACKEND_MODEL_KEYS[backend.service];
    if (modelKey) {
      const backendModel = await env.getValue(modelKey);
      if (backendModel) {
        model = backendModel;
      }
    }

    config = {
      provider,
      model,
      baseUrl: ensureV1Suffix(backend.info.url),
      apiKeyEnv: deriveApiKeyEnv(provider),
      apiKey: userApiKey,
      decayRate,
      startupTimeout,
      localBackend: true,
    };

    addBackendDependency(svc, backend.service);
  } else {
    config = {
      provider: userProvider || 'ollama',
      model: userModel,
      baseUrl: '',
      apiKeyEnv: deriveApiKeyEnv(userProvider || 'ollama'),
      apiKey: userApiKey,
      decayRate,
      startupTimeout,
      localBackend: false,
    };
  }

  const configToml = generateConfigToml(config);
  const script = generateEntrypointScript(configToml, config);

  svc.entrypoint = ['/bin/bash', '-c', script, '--'];

  const envVars: Record<string, string> = {
    HARBOR_OPENFANG_MODEL_PROVIDER: config.provider,
    HARBOR_OPENFANG_MODEL: config.model,
    HARBOR_OPENFANG_BASE_URL: config.baseUrl,
  };

  if (config.apiKeyEnv) {
    envVars[config.apiKeyEnv] = DUMMY_API_KEY;
  }

  setEnvVars(svc, envVars);

  return compose;
}
