# expand_knowledge_base.py

import os
import yaml
import logging
import time
import openai
import argparse
from config import Config

# --- Konfiguration ---
BASE_KNOWLEDGE_FILE = "marketing_wissen.yaml"
OUTPUT_FILE = "marketing_wissen_komplett.yaml"
MODEL_TO_USE = "gpt-4o"

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def call_openai_with_retry(prompt, is_extraction=False, max_retries=3, delay=5):
    # ... (Diese Funktion bleibt unverändert) ...
    for attempt in range(max_retries):
        try:
            logging.info(f"Sende Prompt an OpenAI (Länge: {len(prompt)} Zeichen)...")
            response_format = {"type": "json_object"} if is_extraction else {"type": "text"}
            response = openai.ChatCompletion.create(
                model=MODEL_TO_USE,
                response_format=response_format,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.3,
                max_tokens=2048
            )
            content = response.choices[0].message['content'].strip()
            return content
        except Exception as e:
            logging.error(f"Fehler bei OpenAI-API-Aufruf: {e}")
            if attempt < max_retries - 1:
                time.sleep(delay)
            else:
                return None

def generate_research_prompt(branch_name):
    # ... (Diese Funktion bleibt unverändert) ...
    return (
        f"Erstelle ein prägnantes Branchen-Dossier (ca. 300-400 Wörter) für: '{branch_name}'.\n"
        "Struktur des Dossiers:\n"
        "1.  **Geschäftsmodelle & Field Service:** Beschreibe kurz die typischen Geschäftsmodelle und die zentrale Rolle des technischen Außendienstes (Field Service) in dieser Branche.\n"
        "2.  **Herausforderungen & Trends:** Nenne die wichtigsten aktuellen Herausforderungen und Trends, die den Service-Bereich beeinflussen (z.B. Digitalisierung, Regularien, Fachkräftemangel).\n"
        "3.  **Branchenspezifisches Wording:** Liste einige typische Fachbegriffe oder Abkürzungen auf, die im Service-Kontext dieser Branche üblich sind."
    )

def generate_extraction_prompt(dossier_content):
    """Erstellt den Prompt, um die strukturierten Daten aus dem Dossier zu extrahieren."""
    return (
        "Du bist ein Branchenanalyst mit dem Spezialgebiet Field Service Management. Deine Aufgabe ist es, aus einem Branchen-Dossier die Kernaussagen zu extrahieren.\n"
        "Gib das Ergebnis ausschließlich als sauberes JSON-Objekt mit den Schlüsseln 'summary', 'pain_points' und 'key_terms' aus.\n\n"
        "WICHTIGE REGELN FÜR 'pain_points':\n"
        "- Extrahiere 5 **operative Schmerzpunkte, die direkt den technischen Außendienst betreffen**.\n"
        "- Formuliere sie als konkrete Probleme, die ein Service-Leiter lösen muss (z.B. 'Sicherstellung der Anlagenverfügbarkeit', 'Lückenlose Dokumentation für Audits').\n"
        "- Vermeide allgemeine Management-Themen wie 'Komplexität der Geschäftsmodelle' oder reine HR-Themen wie 'Fachkräftemangel'.\n\n"
        "--- DOSSIER ---\n"
        f"{dossier_content}"
    )

def main(branches_to_process=None):
    """Erweitert die Wissensbasis um die fehlenden Branchen und speichert die Recherche-Dossiers."""
    logging.info("Starte Erweiterung der Wissensbasis...")
    
    Config.load_api_keys()
    openai.api_key = Config.API_KEYS.get('openai')
    if not openai.api_key:
        logging.critical("OpenAI API Key nicht gefunden.")
        return

    try:
        with open(BASE_KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
            knowledge_base = yaml.safe_load(f)
    except FileNotFoundError:
        logging.critical(f"FEHLER: Basis-Wissensdatei '{BASE_KNOWLEDGE_FILE}' nicht gefunden.")
        return

    all_branches = set(Config.BRANCH_GROUP_MAPPING.keys())
    existing_branches = set(knowledge_base.get('Branchen', {}).keys())
    
    if branches_to_process:
        target_branches = [b for b in branches_to_process if b in all_branches]
        if not target_branches:
            logging.error("Keine der angegebenen Branchen ist gültig. Bitte prüfen Sie die Schreibweise.")
            logging.info(f"Gültige Branchen sind: {list(all_branches)}")
            return
        logging.info(f"Verarbeite die {len(target_branches)} explizit angegebenen Branchen...")
    else:
        target_branches = sorted(list(all_branches - existing_branches))
        if not target_branches:
            logging.info("Glückwunsch! Alle Branchen sind bereits in der Wissensbasis vorhanden.")
            return
        logging.info(f"Es werden {len(target_branches)} fehlende Branchen verarbeitet...")

    logging.info(f"Zu verarbeitende Branchen: {', '.join(target_branches)}")
    
    # KORRIGIERTE ZEILE
    DOSSIER_FOLDER = "industries"
    os.makedirs(DOSSIER_FOLDER, exist_ok=True)

    for branch in target_branches:
        if not branches_to_process and branch in existing_branches:
            logging.debug(f"Branche '{branch}' bereits vorhanden, wird übersprungen.")
            continue
            
        logging.info(f"\n--- Verarbeite Branche: {branch} ---")
        
        logging.info("  -> Stufe 1: Generiere Recherche-Dossier...")
        research_prompt = generate_research_prompt(branch)
        dossier = call_openai_with_retry(research_prompt)
        if not dossier: continue
        
        try:
            sanitized_branch_name = branch.replace('/', '-').replace('\\', '-')
            dossier_filepath = os.path.join(DOSSIER_FOLDER, f"{sanitized_branch_name}.txt")
            with open(dossier_filepath, 'w', encoding='utf-8') as f: f.write(dossier)
            logging.info(f"    -> Dossier erfolgreich in '{dossier_filepath}' gespeichert.")
        except Exception as e:
            logging.error(f"    -> Fehler beim Speichern des Dossiers für {branch}: {e}")
        
        time.sleep(2)

        logging.info("  -> Stufe 2: Extrahiere strukturierte Daten aus dem Dossier...")
        extraction_prompt = generate_extraction_prompt(dossier)
        extracted_data_str = call_openai_with_retry(extraction_prompt, is_extraction=True)
        if not extracted_data_str: continue
        
        try:
            if extracted_data_str.startswith("```"):
                extracted_data_str = extracted_data_str.split('\n', 1)[1].rsplit('```', 1)[0]
            
            extracted_data = yaml.safe_load(extracted_data_str)
            extracted_data['references_DE'] = '[HIER DEUTSCHE REFERENZKUNDEN EINTRAGEN]'
            extracted_data['references_GB'] = '[HIER ENGLISCHE REFERENZKUNDEN EINTRAGEN]'
            knowledge_base['Branchen'][branch] = extracted_data
            logging.info(f"    -> {branch} erfolgreich zur Wissensbasis hinzugefügt.")
        except Exception as e:
            logging.error(f"    -> Fehler beim Parsen der extrahierten Daten für {branch}: {e}")
            
        time.sleep(2)

    try:
        with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
            yaml.dump(knowledge_base, f, allow_unicode=True, sort_keys=False, width=120)
        logging.info(f"\nErfolgreich! Die aktualisierte Wissensbasis wurde in '{OUTPUT_FILE}' gespeichert.")
    except Exception as e:
        logging.error(f"Fehler beim Speichern der finalen YAML-Datei: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Erweitert die Marketing-Wissensbasis um fehlende Branchen.")
    parser.add_argument(
        "--branches", 
        nargs='+',
        type=str, 
        help="Eine oder mehrere spezifische Branchen, die verarbeitet werden sollen. Bei Angabe werden nur diese bearbeitet."
    )
    args = parser.parse_args()
    
    main(branches_to_process=args.branches)