import re

# 1. Update types.ts
with open('lib/ai/types.ts', 'r') as f:
    types_content = f.read()

types_content = types_content.replace(
    'generateTags(content: string): Promise<TagSuggestion[]>',
    'generateTags(content: string, language?: string): Promise<TagSuggestion[]>'
)
with open('lib/ai/types.ts', 'w') as f:
    f.write(types_content)

# 2. Update OllamaProvider
with open('lib/ai/providers/ollama.ts', 'r') as f:
    ollama_content = f.read()

ollama_content = ollama_content.replace(
    'async generateTags(content: string): Promise<TagSuggestion[]>',
    'async generateTags(content: string, language: string = "en"): Promise<TagSuggestion[]>'
)

# Replace the hardcoded prompt build logic
prompt_logic = """
      const promptText = language === 'fa' 
        ? `متن زیر را تحلیل کن و مفاهیم کلیدی را به عنوان برچسب استخراج کن (حداکثر ۱-۳ کلمه).\nقوانین:\n- کلمات ربط را حذف کن.\n- عبارات ترکیبی را حفظ کن.\n- حداکثر ۵ برچسب.\nپاسخ فقط به صورت لیست JSON با فرمت [{"tag": "string", "confidence": number}]\nمتن: "${content}"`
        : language === 'fr'
        ? `Analyse la note suivante et extrais les concepts clés sous forme de tags courts (1-3 mots max).\nRègles:\n- Pas de mots de liaison.\n- Garde les expressions composées ensemble.\n- Normalise en minuscules sauf noms propres.\n- Maximum 5 tags.\nRéponds UNIQUEMENT sous forme de liste JSON d'objets : [{"tag": "string", "confidence": number}].\nContenu de la note: "${content}"`
        : `Analyze the following note and extract key concepts as short tags (1-3 words max).\nRules:\n- No stop words.\n- Keep compound expressions together.\n- Lowercase unless proper noun.\n- Max 5 tags.\nRespond ONLY as a JSON list of objects: [{"tag": "string", "confidence": number}].\nNote content: "${content}"`;

      const response = await fetch(`${this.baseUrl}/generate`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({
          model: this.modelName,
          prompt: promptText,
          stream: false,
        }),
      });
"""

# The original has:
# const response = await fetch(`${this.baseUrl}/generate`, {
#   method: 'POST',
#   headers: { 'Content-Type': 'application/json' },
#   body: JSON.stringify({
#     model: this.modelName,
#     prompt: `Analyse la note suivante...

ollama_content = re.sub(
    r'const response = await fetch\(`\$\{this\.baseUrl\}/generate`.*?\}\),\n\s*\}\);',
    prompt_logic.strip(),
    ollama_content,
    flags=re.DOTALL
)

with open('lib/ai/providers/ollama.ts', 'w') as f:
    f.write(ollama_content)

# 3. Update route.ts
with open('app/api/ai/tags/route.ts', 'r') as f:
    route_content = f.read()

route_content = route_content.replace(
    'const tags = await provider.generateTags(content);',
    'const tags = await provider.generateTags(content, language);'
)
with open('app/api/ai/tags/route.ts', 'w') as f:
    f.write(route_content)