import { z } from 'zod';
import { env } from '../../config/env.js';
import type { Stack, Level, QuestionType } from '../../db/schema/enums.js';

export interface LlmConfig {
  baseUrl: string;
  model: string;
  fallbackModel?: string;
  apiKey?: string;
  timeoutMs: number;
  temperature: number;
  maxTokens: number;
  maxRetries: number;
  retryDelayMs: number;
}

export interface ChatMessage {
  role: 'system' | 'user' | 'assistant';
  content: string;
}

export interface ChatCompletionResponse {
  choices: Array<{
    message?: { content: string };
    text?: string;
  }>;
}

const QUESTION_TYPES: QuestionType[] = ['single_choice', 'multiple_select', 'true_false', 'short_text'];

const optionSchema = z.object({
  key: z.string().min(1),
  text: z.string().min(1),
});

const generatedQuestionSchema = z.object({
  questionText: z.string().min(1),
  type: z.enum(QUESTION_TYPES as [string, ...string[]]),
  options: z.array(optionSchema).optional(),
  correctAnswer: z.union([z.string(), z.array(z.string())]),
  explanation: z.string().min(1),
});

const generateQuestionsResponseSchema = z.object({
  questions: z.array(generatedQuestionSchema),
});

export type GeneratedQuestion = z.infer<typeof generatedQuestionSchema> & {
  stack: Stack;
  level: Level;
};

export interface GenerateQuestionsInput {
  stack: Stack;
  level: Level;
  count: number;
  types?: QuestionType[];
}

export class LlmService {
  private readonly config: LlmConfig;

  constructor(config?: Partial<LlmConfig>) {
    this.config = {
      baseUrl: config?.baseUrl ?? env.LLM_BASE_URL,
      model: config?.model ?? env.LLM_MODEL,
      fallbackModel: config?.fallbackModel ?? env.LLM_FALLBACK_MODEL,
      apiKey: config?.apiKey ?? env.LLM_API_KEY,
      timeoutMs: config?.timeoutMs ?? env.LLM_TIMEOUT_MS,
      temperature: config?.temperature ?? env.LLM_TEMPERATURE,
      maxTokens: config?.maxTokens ?? env.LLM_MAX_TOKENS,
      maxRetries: config?.maxRetries ?? env.LLM_MAX_RETRIES,
      retryDelayMs: config?.retryDelayMs ?? env.LLM_RETRY_DELAY_MS,
    };
  }

  async chat(messages: ChatMessage[]): Promise<string> {
    let lastError: Error | null = null;

    const modelsToTry = [this.config.model];
    if (this.config.fallbackModel) {
      modelsToTry.push(this.config.fallbackModel);
    }

    for (const model of modelsToTry) {
      for (let attempt = 0; attempt <= this.config.maxRetries; attempt++) {
        try {
          return await this.executeChat(messages, model);
        } catch (err) {
          lastError = err instanceof Error ? err : new Error('LLM request failed');
          if (attempt < this.config.maxRetries) {
            const delayMs = this.config.retryDelayMs * Math.pow(2, attempt);
            await sleep(delayMs);
          }
        }
      }
    }

    throw lastError ?? new Error('LLM request failed');
  }

  private async executeChat(messages: ChatMessage[], model: string): Promise<string> {
    const url = `${this.config.baseUrl.replace(/\/$/, '')}/chat/completions`;

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };
    if (this.config.apiKey) {
      headers['Authorization'] = `Bearer ${this.config.apiKey}`;
    }

    const body = {
      model,
      messages: messages.map((m) => ({ role: m.role, content: m.content })),
      temperature: this.config.temperature,
      max_tokens: this.config.maxTokens,
    };

    const controller = new AbortController();
    const timeoutId = setTimeout(() => controller.abort(), this.config.timeoutMs);

    try {
      const res = await fetch(url, {
        method: 'POST',
        headers,
        body: JSON.stringify(body),
        signal: controller.signal,
      });

      clearTimeout(timeoutId);

      if (!res.ok) {
        const text = await res.text();
        throw new Error(`LLM request failed: ${res.status} ${res.statusText} - ${text}`);
      }

      const data = (await res.json()) as ChatCompletionResponse;

      const choice = data.choices?.[0];
      const content = choice?.message?.content ?? choice?.text ?? '';

      return content.trim();
    } catch (err) {
      clearTimeout(timeoutId);
      if (err instanceof Error) {
        throw err;
      }
      throw new Error('LLM request failed');
    }
  }

  async generateQuestions(input: GenerateQuestionsInput): Promise<GeneratedQuestion[]> {
    const { stack, level, count, types = QUESTION_TYPES } = input;

    const typeList = types.join(', ');
    const systemPrompt = `You are a technical interview question generator. Generate exactly ${count} programming/tech questions.
Return ONLY valid JSON in this exact format (no markdown, no code blocks):
{"questions":[{"questionText":"...","type":"single_choice|multiple_select|true_false|short_text","options":[{"key":"a","text":"..."}],"correctAnswer":"a" or ["a","b"],"explanation":"..."}]}
Rules: type must be one of: ${typeList}. For single_choice/multiple_select: options array required with key (a,b,c,d). For true_false: options [{"key":"true","text":"True"},{"key":"false","text":"False"}]. For short_text: options omitted, correctAnswer is string.`;

    const userPrompt = `Generate ${count} questions for stack="${stack}", level="${level}". Use types: ${typeList}.`;

    const raw = await this.chat([
      { role: 'system', content: systemPrompt },
      { role: 'user', content: userPrompt },
    ]);

    const jsonStr = extractJson(raw);
    const parsed = JSON.parse(jsonStr) as unknown;

    const result = generateQuestionsResponseSchema.safeParse(parsed);
    if (!result.success) {
      throw new Error(`LLM response validation failed: ${result.error.message}`);
    }

    const questions: GeneratedQuestion[] = result.data.questions.map((q) => ({
      ...q,
      stack,
      level,
    }));

    for (const q of questions) {
      if ((q.type === 'single_choice' || q.type === 'multiple_select') && (!q.options || q.options.length === 0)) {
        throw new Error(`Question validation failed: ${q.type} requires options`);
      }
      if (q.type === 'true_false' && (!q.options || q.options.length < 2)) {
        throw new Error(`Question validation failed: true_false requires true/false options`);
      }
    }

    return questions;
  }
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function extractJson(text: string): string {
  const trimmed = text.trim();
  const match = trimmed.match(/\{[\s\S]*\}/);
  return match ? match[0]! : trimmed;
}