import { z } from 'zod'

const ProviderSchema = z.enum(['openrouter', 'deepseek'])
const CompletionResponseSchema = z.object({
  choices: z.array(z.object({
    message: z.object({ content: z.string().min(1) }),
  })).min(1),
  usage: z.object({
    prompt_tokens: z.number().nonnegative().nullish(),
    completion_tokens: z.number().nonnegative().nullish(),
    input_tokens: z.number().nonnegative().nullish(),
    output_tokens: z.number().nonnegative().nullish(),
    cost: z.number().nonnegative().nullish(),
    total_cost: z.number().nonnegative().nullish(),
  }).optional(),
})

export type AiProvider = z.infer<typeof ProviderSchema>

export interface AiRuntimeConfig {
  aiProvider?: unknown
  openrouterApiKey?: unknown
  openrouterModel?: unknown
  openrouterApiEndpoint?: unknown
  deepseekApiKey?: unknown
  deepseekModel?: unknown
  deepseekApiEndpoint?: unknown
}

export interface JsonCompletionRequest {
  system: string
  messages: Array<{ role: 'user' | 'assistant'; content: string }>
  schemaName: string
  jsonSchema: Record<string, unknown>
}

export interface JsonCompletion {
  endpoint: string
  headers: Record<string, string>
  body: Record<string, unknown>
  provider: AiProvider
  model: string
}

export interface AiUsageSnapshot {
  inputTokens: number
  outputTokens: number
  costUsd: number
}

const nonEmptyString = (value: unknown, name: string) => {
  const result = z.string().trim().min(1).safeParse(value)
  if (!result.success) throw new Error(`${name} is required for the selected AI provider`)
  return result.data
}

const endpoint = (value: unknown, fallback: string, name: string) => {
  const result = z.string().url().safeParse(value || fallback)
  if (!result.success) throw new Error(`${name} must be a valid URL`)
  return result.data
}

export function resolveAiProvider(config: AiRuntimeConfig) {
  const providerResult = ProviderSchema.safeParse(config.aiProvider || 'openrouter')
  if (!providerResult.success) throw new Error('AI_PROVIDER must be either "openrouter" or "deepseek"')

  if (providerResult.data === 'deepseek') {
    return {
      provider: 'deepseek' as const,
      apiKey: nonEmptyString(config.deepseekApiKey, 'DEEPSEEK_API_KEY'),
      model: nonEmptyString(config.deepseekModel || 'deepseek-v4-flash', 'DEEPSEEK_MODEL'),
      endpoint: endpoint(config.deepseekApiEndpoint, 'https://api.deepseek.com/chat/completions', 'DEEPSEEK_API_ENDPOINT'),
    }
  }

  return {
    provider: 'openrouter' as const,
    apiKey: nonEmptyString(config.openrouterApiKey, 'OPENROUTER_API_KEY'),
    model: nonEmptyString(config.openrouterModel || 'deepseek/deepseek-v4-flash', 'OPENROUTER_MODEL'),
    endpoint: endpoint(config.openrouterApiEndpoint, 'https://openrouter.ai/api/v1/chat/completions', 'OPENROUTER_API_ENDPOINT'),
  }
}

export function buildJsonCompletion(config: AiRuntimeConfig, request: JsonCompletionRequest): JsonCompletion {
  const provider = resolveAiProvider(config)
  const messages = [{ role: 'system' as const, content: request.system }, ...request.messages]
  const common = { model: provider.model, messages, stream: false }

  if (provider.provider === 'deepseek') {
    return {
      endpoint: provider.endpoint,
      provider: provider.provider,
      model: provider.model,
      headers: { Authorization: `Bearer ${provider.apiKey}`, 'Content-Type': 'application/json' },
      body: {
        ...common,
        messages: [
          { role: 'system' as const, content: `${request.system}\nReturn only valid JSON matching this JSON Schema: ${JSON.stringify(request.jsonSchema)}` },
          ...request.messages,
        ],
        response_format: { type: 'json_object' },
        thinking: { type: 'disabled' },
      },
    }
  }

  return {
    endpoint: provider.endpoint,
    provider: provider.provider,
    model: provider.model,
    headers: { Authorization: `Bearer ${provider.apiKey}`, 'Content-Type': 'application/json', 'X-Title': 'Dungeons & Ground' },
    body: {
      ...common,
      response_format: { type: 'json_schema', json_schema: { name: request.schemaName, strict: true, schema: request.jsonSchema } },
      provider: { require_parameters: true, data_collection: 'deny' },
    },
  }
}

export function parseJsonCompletion(payload: unknown): unknown {
  const completion = CompletionResponseSchema.parse(payload)
  return JSON.parse(completion.choices[0]!.message.content)
}

export function parseAiUsage(payload: unknown): AiUsageSnapshot {
  const completion = CompletionResponseSchema.parse(payload)
  return {
    inputTokens: Math.max(0, Math.trunc(completion.usage?.prompt_tokens ?? completion.usage?.input_tokens ?? 0)),
    outputTokens: Math.max(0, Math.trunc(completion.usage?.completion_tokens ?? completion.usage?.output_tokens ?? 0)),
    costUsd: Math.max(0, completion.usage?.cost ?? completion.usage?.total_cost ?? 0),
  }
}