import { SearchResult, ChatMessage } from "../types";

export interface PromptTemplate {
  id: string;
  name: string;
  template: string;
  variables: string[];
  description?: string;
  language?: string;
}

export interface PromptContext {
  query: string;
  retrievedDocuments: SearchResult[];
  conversationHistory: ChatMessage[];
  userContext?: Record<string, any>;
  language?: string;
}

export class PromptManager {
  private templates: Map<string, PromptTemplate> = new Map();

  constructor() {
    this.initializeDefaultTemplates();
  }

  private initializeDefaultTemplates(): void {
    // English RAG prompt template
    this.addTemplate({
      id: "rag-default-en",
      name: "Default RAG Prompt (English)",
      template: `You are an AI assistant that answers questions based on the provided context documents. Follow these guidelines:

1. Use only information from the provided context documents
2. If the context doesn't contain enough information to answer the question, say so clearly
3. Cite specific sources when possible using [Source: document_name] format
4. Be concise but comprehensive in your answers
5. If there are conflicting information in the sources, mention this

Context Documents:
{context}

Conversation History:
{conversation_history}

Question: {query}

Answer:`,
      variables: ["context", "conversation_history", "query"],
      description: "Standard RAG prompt for English conversations",
      language: "en",
    });

    // Korean RAG prompt template
    this.addTemplate({
      id: "rag-default-ko",
      name: "Default RAG Prompt (Korean)",
      template: `당신은 제공된 문서 내용을 바탕으로 질문에 답하는 AI 어시스턴트입니다. 다음 지침을 따르세요:

1. 제공된 문서의 정보만을 사용하여 답변하세요
2. 문서에 충분한 정보가 없으면 명확히 알려주세요
3. 가능한 경우 [출처: 문서명] 형식으로 출처를 명시하세요
4. 간결하면서도 포괄적인 답변을 제공하세요
5. 출처들 간에 상충하는 정보가 있으면 이를 언급하세요

문서 내용:
{context}

대화 기록:
{conversation_history}

질문: {query}

답변:`,
      variables: ["context", "conversation_history", "query"],
      description: "Korean용 표준 RAG 프롬프트",
      language: "ko",
    });

    // Conversational RAG prompt
    this.addTemplate({
      id: "rag-conversational-en",
      name: "Conversational RAG Prompt",
      template: `You are a helpful AI assistant having a conversation with a user. You have access to a knowledge base through document retrieval.

Instructions:
- Use the retrieved documents to inform your responses
- Maintain conversational context from previous messages
- Be natural and engaging while staying factual
- If documents don't contain relevant information, use your general knowledge but mention the limitation
- Reference sources when making specific claims: [Source: document_name]

Retrieved Documents:
{context}

Previous Conversation:
{conversation_history}

User: {query}
Assistant:`,
      variables: ["context", "conversation_history", "query"],
      description:
        "More conversational RAG prompt that allows general knowledge",
      language: "en",
    });

    // Summary prompt
    this.addTemplate({
      id: "summarize-documents",
      name: "Document Summarization",
      template: `Please provide a comprehensive summary of the following documents:

{context}

Requirements:
- Identify key themes and main points
- Maintain factual accuracy
- Structure the summary logically
- Note any important details or conclusions
- If documents cover different topics, organize by topic

Summary:`,
      variables: ["context"],
      description: "Template for summarizing retrieved documents",
    });

    // Question generation prompt
    this.addTemplate({
      id: "generate-questions",
      name: "Question Generation",
      template: `Based on the following documents, generate relevant follow-up questions that users might ask:

{context}

Generate 3-5 specific, actionable questions that would help users explore this topic further. Focus on:
- Practical applications
- Deeper understanding
- Related concepts
- Implementation details

Questions:`,
      variables: ["context"],
      description: "Generate follow-up questions based on context",
    });
  }

  addTemplate(template: PromptTemplate): void {
    this.templates.set(template.id, template);
  }

  getTemplate(id: string): PromptTemplate | null {
    return this.templates.get(id) || null;
  }

  listTemplates(): PromptTemplate[] {
    return Array.from(this.templates.values());
  }

  buildPrompt(templateId: string, context: PromptContext): string {
    const template = this.getTemplate(templateId);
    if (!template) {
      throw new Error(`Template not found: ${templateId}`);
    }

    let prompt = template.template;

    // Replace context
    const contextText = this.formatContext(context.retrievedDocuments);
    prompt = prompt.replace("{context}", contextText);

    // Replace conversation history
    const historyText = this.formatConversationHistory(
      context.conversationHistory
    );
    prompt = prompt.replace("{conversation_history}", historyText);

    // Replace query
    prompt = prompt.replace("{query}", context.query);

    // Replace any custom variables from userContext
    if (context.userContext) {
      for (const [key, value] of Object.entries(context.userContext)) {
        const placeholder = `{${key}}`;
        prompt = prompt.replace(new RegExp(placeholder, "g"), String(value));
      }
    }

    return prompt;
  }

  private formatContext(documents: SearchResult[]): string {
    if (documents.length === 0) {
      return "No relevant documents found.";
    }

    return documents
      .map((result, index) => {
        const docTitle =
          result.document.metadata.title || `Document ${index + 1}`;
        const chunkContent = result.chunk.content.trim();
        const score = (result.score * 100).toFixed(1);

        return `[Document ${index + 1}: ${docTitle}] (Relevance: ${score}%)
${chunkContent}`;
      })
      .join("\n\n---\n\n");
  }

  private formatConversationHistory(messages: ChatMessage[]): string {
    if (messages.length === 0) {
      return "No previous conversation.";
    }

    // Filter out system messages and format user/assistant messages
    const conversationMessages = messages.filter(
      (msg) => msg.role !== "system"
    );

    if (conversationMessages.length === 0) {
      return "No previous conversation.";
    }

    return conversationMessages
      .map((msg) => {
        const role = msg.role === "user" ? "User" : "Assistant";
        return `${role}: ${msg.content}`;
      })
      .join("\n");
  }

  // Get template by language preference
  getTemplateByLanguage(
    baseId: string,
    language: string
  ): PromptTemplate | null {
    const languageSpecificId = `${baseId}-${language}`;
    return this.getTemplate(languageSpecificId) || this.getTemplate(baseId);
  }

  // Validate template variables
  validateTemplate(template: PromptTemplate): {
    isValid: boolean;
    missingVariables: string[];
  } {
    const requiredVariables = ["context", "query"]; // Basic required variables
    const templateText = template.template;
    const missingVariables: string[] = [];

    for (const variable of requiredVariables) {
      const placeholder = `{${variable}}`;
      if (!templateText.includes(placeholder)) {
        missingVariables.push(variable);
      }
    }

    return {
      isValid: missingVariables.length === 0,
      missingVariables,
    };
  }

  // Dynamic template creation
  createCustomTemplate(
    id: string,
    name: string,
    template: string,
    options: {
      description?: string;
      language?: string;
      variables?: string[];
    } = {}
  ): PromptTemplate {
    // Extract variables from template
    const variableMatches = template.match(/{([^}]+)}/g) || [];
    const extractedVariables = variableMatches.map((match) =>
      match.slice(1, -1)
    );

    const promptTemplate: PromptTemplate = {
      id,
      name,
      template,
      variables: options.variables || extractedVariables,
      description: options.description,
      language: options.language,
    };

    this.addTemplate(promptTemplate);
    return promptTemplate;
  }

  // Remove template
  removeTemplate(id: string): boolean {
    return this.templates.delete(id);
  }

  // Get template statistics
  getTemplateStats(): {
    totalTemplates: number;
    languageBreakdown: Record<string, number>;
    averageVariables: number;
  } {
    const templates = this.listTemplates();
    const languageBreakdown: Record<string, number> = {};
    let totalVariables = 0;

    for (const template of templates) {
      const language = template.language || "unknown";
      languageBreakdown[language] = (languageBreakdown[language] || 0) + 1;
      totalVariables += template.variables.length;
    }

    return {
      totalTemplates: templates.length,
      languageBreakdown,
      averageVariables:
        templates.length > 0 ? totalVariables / templates.length : 0,
    };
  }
}