#!/usr/bin/env python3
import os
import sys
import time
import random
import logging
import requests
import json
import re
import subprocess
import shutil
from datetime import datetime
from pathlib import Path
import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from bs4 import BeautifulSoup
import html2text

# Importation du module de configuration
from config import app_config as config

# Configuration de la page Streamlit
st.set_page_config(
    page_title="DevSecOps Data Bot",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialisation des variables de session
config.init_session_state()

# Configuration du logging
def setup_logging():
    log_dir = Path("logs")
    log_dir.mkdir(exist_ok=True)
    log_file = log_dir / f"data_collector_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler(sys.stdout)
        ]
    )
    return logging.getLogger(__name__)

logger = setup_logging()

# Création des répertoires et scripts nécessaires
def create_initial_setup():
    dirs = [
        "data/devsecops/qa", "data/security/qa", "data/development/qa",
        "data/data_analysis/qa", "logs", "config", "server", "scripts",
        "models", "llama.cpp", ".kaggle"
    ]
    for dir_path in dirs:
        Path(dir_path).mkdir(parents=True, exist_ok=True)
    
    download_script = Path("scripts/download_with_aria2c.sh")
    if not download_script.exists():
        with open(download_script, 'w') as f:
            f.write("""#!/bin/bash
URL=$1
OUTPUT=$2
MAX_RETRIES=5
for i in $(seq 1 $MAX_RETRIES); do
    echo "Tentative $i/$MAX_RETRIES: $URL"
    aria2c -x 16 -s 16 -d "$(dirname "$OUTPUT")" -o "$(basename "$OUTPUT")" "$URL" && break
    sleep 10
done
""")
        os.chmod(download_script, 0o755)
    
    llama_dir = Path("llama.cpp")
    if not llama_dir.exists():
        st.info("Installation de llama.cpp...")
        subprocess.run(["git", "clone", "https://github.com/ggerganov/llama.cpp.git", str(llama_dir)])
        os.chdir(str(llama_dir))
        subprocess.run(["mkdir", "-p", "build"])
        os.chdir("build")
        subprocess.run(["cmake", "..", "-DLLAMA_CURL=1"])
        subprocess.run(["cmake", "--build", ".", "--config", "Release"])
        os.chdir(Path(__file__).parent)

# Convertisseur HTML vers texte
h = html2text.HTML2Text()
h.ignore_links = False
h.ignore_images = True
h.ignore_emphasis = False
h.body_width = 0

# Fonctions pour le serveur LLM (llama.cpp)
def check_server_status():
    try:
        response = requests.get(config.LLM_SERVER_URL.replace("/completion", "/health"), timeout=2)
        if response.status_code == 200:
            st.session_state.server_status = "Actif"
            return True
        else:
            st.session_state.server_status = "Inactif"
            return False
    except requests.exceptions.RequestException:
        st.session_state.server_status = "Inactif"
        return False

def start_llm_server():
    if check_server_status():
        st.info("Le serveur llama.cpp est déjà en cours d'exécution.")
        return
    
    model_path = Path("models/qwen2.5-coder-1.5b-q8_0.gguf")
    if not model_path.exists():
        st.error("Le modèle GGUF n'existe pas. Veuillez le placer dans le dossier models/.")
        return
    
    llama_server = Path("llama.cpp/build/bin/llama-server")
    if not llama_server.exists():
        st.error("llama.cpp n'est pas compilé. Veuillez compiler llama.cpp d'abord.")
        return
    
    start_script = Path("server/start_server.sh")
    if not start_script.exists():
        with open(start_script, 'w') as f:
            f.write(f"""#!/bin/bash
MODEL_PATH="{str(model_path)}"
if [ ! -f "$MODEL_PATH" ]; then
    echo "Le modèle GGUF est introuvable à: $MODEL_PATH"
    exit 1
fi
"{str(llama_server)}" \\
  -m "$MODEL_PATH" \\
  --port 8080 \\
  --host 0.0.0.0 \\
  -c 4096 \\
  -ngl 999 \\
  --threads 8 \\
  > "logs/llama_server.log" 2>&1 &
echo $! > "server/server.pid"
""")
        os.chmod(start_script, 0o755)
    
    try:
        subprocess.Popen(["bash", str(start_script)])
        st.success("Le serveur llama.cpp est en cours de démarrage...")
        time.sleep(5)
        if check_server_status():
            st.success("Serveur llama.cpp démarré avec succès!")
        else:
            st.error("Le serveur n'a pas pu démarrer. Vérifiez les logs dans le dossier logs/.")
    except Exception as e:
        st.error(f"Erreur lors du démarrage du serveur: {str(e)}")

def stop_llm_server():
    stop_script = Path("server/stop_server.sh")
    if not stop_script.exists():
        with open(stop_script, 'w') as f:
            f.write("""#!/bin/bash
PID_FILE="server/server.pid"
if [ -f "$PID_FILE" ]; then
    PID=$(cat "$PID_FILE")
    kill $PID
    rm "$PID_FILE"
    echo "Serveur llama.cpp arrêté."
else
    echo "Aucun PID de serveur trouvé."
fi
""")
        os.chmod(stop_script, 0o755)
    
    try:
        subprocess.run(["bash", str(stop_script)])
        st.success("Le serveur llama.cpp est en cours d'arrêt...")
        time.sleep(2)
        if not check_server_status():
            st.success("Serveur llama.cpp arrêté avec succès!")
        else:
            st.warning("Le serveur n'a pas pu être arrêté correctement.")
    except Exception as e:
        st.error(f"Erreur lors de l'arrêt du serveur: {str(e)}")

def load_prompts():
    prompts_file = Path("config/prompts.json")
    if not prompts_file.exists():
        default_prompts = {
            "enrich_qa": {
                "system": "Tu es un expert DevSecOps. Améliore cette paire question/réponse en y ajoutant des tags, des signatures d'attaques potentielles, et en structurant les informations. Réponds uniquement avec un objet JSON.",
                "prompt_template": "Question originale: {question}\nRéponse originale: {answer}\nContexte: {context}\n\nFournis une version améliorée sous forme de JSON:\n{{\n  \"question\": \"question améliorée\",\n  \"answer\": \"réponse améliorée\",\n  \"tags\": [\"tag1\", \"tag2\"],\n  \"attack_signatures\": [\"signature1\", \"signature2\"]\n}}"
            },
            "analyze_relevance": {
                "system": "Analyse ce contenu et détermine s'il est pertinent pour DevSecOps. Si pertinent, extrais les signatures d'attaques connues. Réponds uniquement avec un objet JSON.",
                "prompt_template": "Contenu: {content}...\n\nRéponds sous forme de JSON:\n{{\n  \"relevant\": true,\n  \"attack_signatures\": [\"signature1\", \"signature2\"],\n  \"security_tags\": [\"tag1\", \"tag2\"],\n  \"it_relevance_score\": 0-100\n}}"
            },
            "generate_queries": {
                "system": "Analyse les données actuelles et génère 5 nouvelles requêtes de recherche pour trouver plus de contenu DevSecOps pertinent, en particulier des signatures d'attaques et vulnérabilités. Réponds uniquement avec un objet JSON.",
                "prompt_template": "Données actuelles: {current_data}...\n\nRéponds sous forme de JSON:\n{{\n  \"queries\": [\"query1\", \"query2\", \"query3\", \"query4\", \"query5\"]\n}}"
            }
        }
        with open(prompts_file, 'w') as f:
            json.dump(default_prompts, f, indent=2)
    
    with open(prompts_file, 'r', encoding='utf-8') as f:
        return json.load(f)

PROMPTS = load_prompts()

class IAEnricher:
    def __init__(self):
        self.server_url = config.LLM_SERVER_URL
        self.available = check_server_status()
        if self.available:
            logger.info("Serveur llama.cpp détecté et prêt.")
        else:
            logger.warning("Serveur llama.cpp non disponible. Les fonctionnalités d'IA seront désactivées.")
    
    def _get_qwen_response(self, prompt, **kwargs):
        if not self.available:
            return None
        
        payload = {
            "prompt": prompt,
            "n_predict": kwargs.get('n_predict', 512),
            "temperature": kwargs.get('temperature', 0.7),
            "stop": ["<|im_end|>", "</s>"]
        }
        
        try:
            response = requests.post(self.server_url, json=payload, timeout=60)
            if response.status_code == 200:
                return response.json().get('content', '')
            else:
                logger.error(f"Erreur du serveur LLM: {response.status_code} - {response.text}")
                return None
        except requests.exceptions.RequestException as e:
            logger.error(f"Erreur de connexion au serveur LLM: {str(e)}")
            return None
    
    def enrich_qa_pair(self, question, answer, context=""):
        if not self.available or not st.session_state.enable_enrichment:
            return question, answer, [], []
        
        prompt_template = PROMPTS.get("enrich_qa", {}).get("prompt_template", "")
        system_prompt = PROMPTS.get("enrich_qa", {}).get("system", "")
        
        full_prompt = f"{system_prompt}\n\n{prompt_template.format(question=question, answer=answer, context=context[:500])}"
        response_text = self._get_qwen_response(full_prompt, n_predict=1024)
        
        if response_text:
            try:
                json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                if json_match:
                    enriched_data = json.loads(json_match.group())
                    return (
                        enriched_data.get('question', question),
                        enriched_data.get('answer', answer),
                        enriched_data.get('tags', []),
                        enriched_data.get('attack_signatures', [])
                    )
            except json.JSONDecodeError as e:
                logger.warning(f"Erreur de décodage JSON de la réponse IA: {e}")
        
        return question, answer, [], []
    
    def analyze_content_relevance(self, content):
        if not self.available or not st.session_state.enable_enrichment:
            return True, [], [], 50
        
        prompt_template = PROMPTS.get("analyze_relevance", {}).get("prompt_template", "")
        system_prompt = PROMPTS.get("analyze_relevance", {}).get("system", "")
        
        full_prompt = f"{system_prompt}\n\n{prompt_template.format(content=content[:1500])}"
        response_text = self._get_qwen_response(full_prompt, n_predict=256, temperature=st.session_state.temperature)
        
        if response_text:
            try:
                json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                if json_match:
                    analysis = json.loads(json_match.group())
                    return (
                        analysis.get('relevant', True),
                        analysis.get('attack_signatures', []),
                        analysis.get('security_tags', []),
                        analysis.get('it_relevance_score', 50)
                    )
            except json.JSONDecodeError as e:
                logger.warning(f"Erreur de décodage JSON de la réponse IA: {e}")
        return True, [], [], 50
    
    def generate_adaptive_queries(self, current_data):
        if not self.available or not st.session_state.enable_enrichment:
            return []
        
        prompt_template = PROMPTS.get("generate_queries", {}).get("prompt_template", "")
        system_prompt = PROMPTS.get("generate_queries", {}).get("system", "")
        
        full_prompt = f"{system_prompt}\n\n{prompt_template.format(current_data=current_data[:1000])}"
        response_text = self._get_qwen_response(full_prompt, n_predict=st.session_state.n_predict)
        
        if response_text:
            try:
                json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
                if json_match:
                    queries_data = json.loads(json_match.group())
                    return queries_data.get('queries', [])
            except json.JSONDecodeError as e:
                logger.warning(f"Erreur de décodage JSON de la réponse IA: {e}")
        return []

ia_enricher = IAEnricher()

def check_api_keys():
    keys = {
        'GITHUB_API_TOKEN': os.getenv('GITHUB_API_TOKEN'),
        'HUGGINGFACE_API_TOKEN': os.getenv('HUGGINGFACE_API_TOKEN'),
        'NVD_API_KEY': os.getenv('NVD_API_KEY'),
        'STACK_EXCHANGE_API_KEY': os.getenv('STACK_EXCHANGE_API_KEY')
    }
    
    valid_keys = {k: v for k, v in keys.items() if v and v != f'your_{k.lower()}_here'}
    
    config.USE_API_KEYS = len(valid_keys) == len(keys)
    if not config.USE_API_KEYS:
        missing = set(keys.keys()) - set(valid_keys.keys())
        logger.warning(f"Clés d'API manquantes ou non configurées: {', '.join(missing)}")
        logger.warning("Le bot fonctionnera en mode dégradé avec des pauses plus longues.")
    else:
        logger.info("Toutes les clés d'API sont configurées.")
    return config.USE_API_KEYS

def make_request(url, headers=None, params=None, is_api_call=True):
    config.REQUEST_COUNT += 1
    
    pause_factor = 1 if config.USE_API_KEYS else 2
    
    if config.REQUEST_COUNT >= config.MAX_REQUESTS_BEFORE_PAUSE:
        pause_time = random.uniform(config.MIN_PAUSE * pause_factor, config.MAX_PAUSE * pause_factor)
        logger.info(f"Pause de {pause_time:.2f} secondes après {config.REQUEST_COUNT} requêtes...")
        time.sleep(pause_time)
        config.REQUEST_COUNT = 0
    
    try:
        response = requests.get(url, headers=headers, params=params, timeout=30)
        
        if response.status_code == 200:
            return response
        elif response.status_code in [401, 403]:
            logger.warning(f"Accès non autorisé à {url}. Vérifiez vos clés d'API.")
            return None
        elif response.status_code == 429:
            retry_after = int(response.headers.get('Retry-After', 10))
            logger.warning(f"Limite de débit atteinte. Pause de {retry_after} secondes...")
            time.sleep(retry_after)
            return make_request(url, headers, params, is_api_call)
        else:
            logger.warning(f"Statut HTTP {response.status_code} pour {url}")
            return None
    except requests.exceptions.RequestException as e:
        logger.error(f"Erreur lors de la requête à {url}: {str(e)}")
        return None

def clean_html(html_content):
    if not html_content:
        return ""
    text = h.handle(html_content)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def save_qa_pair(question, answer, category, subcategory, source, attack_signatures=None, tags=None):
    if ia_enricher.available and st.session_state.enable_enrichment:
        enriched_question, enriched_answer, enriched_tags, enriched_signatures = ia_enricher.enrich_qa_pair(
            question, answer, f"{category}/{subcategory}"
        )
        
        question = enriched_question
        answer = enriched_answer
        tags = list(set((tags or []) + enriched_tags))
        attack_signatures = list(set((attack_signatures or []) + enriched_signatures))
    
    save_dir = Path("data") / category / "qa"
    save_dir.mkdir(parents=True, exist_ok=True)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"{subcategory}_{source}_{st.session_state.total_qa_pairs}_{timestamp}.json"
    filename = re.sub(r'[^\w\s-]', '', filename).replace(' ', '_')
    
    qa_data = {
        "question": question,
        "answer": answer,
        "category": category,
        "subcategory": subcategory,
        "source": source,
        "timestamp": timestamp,
        "attack_signatures": attack_signatures or [],
        "tags": tags or []
    }
    
    try:
        with open(save_dir / filename, "w", encoding="utf-8") as f:
            json.dump(qa_data, f, indent=2, ensure_ascii=False)
        
        st.session_state.total_qa_pairs += 1
        st.session_state.qa_data.append(qa_data)
        
        logger.info(f"Paire Q/R sauvegardée: {filename} (Total: {st.session_state.total_qa_pairs})")
        st.session_state.logs.append(f"Sauvegardé: {filename}")
    except Exception as e:
        logger.error(f"Erreur lors de la sauvegarde du fichier {filename}: {str(e)}")

def collect_kaggle_data(queries):
    logger.info("Début de la collecte des données Kaggle...")
    
    # Importation de Kaggle ici, après avoir configuré les variables d'environnement
    os.environ['KAGGLE_USERNAME'] = os.getenv('KAGGLE_USERNAME')
    os.environ['KAGGLE_KEY'] = os.getenv('KAGGLE_KEY')
    import kaggle
    kaggle.api.authenticate()
    
    search_queries = queries.split('\n') if queries else ["cybersecurity", "vulnerability"]
    
    if ia_enricher.available and st.session_state.enable_enrichment:
        adaptive_queries = ia_enricher.generate_adaptive_queries("Initial data keywords: " + ", ".join(search_queries))
        search_queries.extend(adaptive_queries)
    
    for query in list(set(search_queries)):
        logger.info(f"Recherche de datasets Kaggle pour: {query}")
        try:
            datasets = kaggle.api.dataset_list(search=query, max_results=5)
            for dataset in datasets:
                dataset_ref = dataset.ref
                if ia_enricher.available and st.session_state.enable_enrichment:
                    is_relevant, _, _, relevance_score = ia_enricher.analyze_content_relevance(dataset.title + " " + dataset.subtitle)
                    if not is_relevant or relevance_score < st.session_state.min_relevance:
                        logger.info(f"Dataset non pertinent ({relevance_score}%): {dataset_ref}. Ignoré.")
                        continue
                
                logger.info(f"Traitement du dataset: {dataset_ref}")
                download_dir = Path("data") / "security" / "kaggle" / dataset_ref.replace('/', '_')
                download_dir.mkdir(parents=True, exist_ok=True)
                kaggle.api.dataset_download_files(dataset_ref, path=download_dir, unzip=True)
                
                for file_path in download_dir.glob('*'):
                    if file_path.is_file() and file_path.suffix.lower() in ['.json', '.csv', '.txt']:
                        try:
                            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                                file_content = f.read()[:5000]
                            is_relevant, signatures, security_tags, _ = ia_enricher.analyze_content_relevance(file_content)
                            if is_relevant:
                                save_qa_pair(
                                    question=f"Quelles informations de sécurité contient le fichier {file_path.name} du dataset '{dataset.title}'?",
                                    answer=file_content, category="security", subcategory="vulnerability",
                                    source=f"kaggle_{dataset_ref}", attack_signatures=signatures, tags=security_tags
                                )
                        except Exception as e:
                            logger.error(f"Erreur lors du traitement du fichier {file_path}: {str(e)}")
                time.sleep(random.uniform(2, 4))
        except Exception as e:
            logger.error(f"Erreur lors de la collecte des données Kaggle pour {query}: {str(e)}")
    logger.info("Collecte des données Kaggle terminée.")

def collect_github_data(queries):
    logger.info("Début de la collecte des données GitHub...")
    base_url = "https://api.github.com"
    headers = {"Accept": "application/vnd.github.v3+json"}
    if config.USE_API_KEYS:
        token = os.getenv('GITHUB_API_TOKEN')
        headers["Authorization"] = f"token {token}"
    
    search_queries = queries.split('\n') if queries else ["topic:devsecops", "topic:security"]
    
    for query in search_queries:
        logger.info(f"Recherche de repositories pour: {query}")
        search_url = f"{base_url}/search/repositories"
        params = {"q": query, "sort": "stars", "per_page": 10}
        response = make_request(search_url, headers=headers, params=params)
        if not response:
            continue
        
        data = response.json()
        for repo in data.get("items", []):
            repo_name = repo["full_name"].replace("/", "_")
            logger.info(f"Traitement du repository: {repo['full_name']}")
            
            issues_url = f"{base_url}/repos/{repo['full_name']}/issues"
            issues_params = {"state": "closed", "labels": "security,bug,vulnerability", "per_page": 10}
            issues_response = make_request(issues_url, headers=headers, params=issues_params)
            
            if issues_response:
                issues_data = issues_response.json()
                for issue in issues_data:
                    if "pull_request" in issue: continue
                    question = issue.get("title", "")
                    body = clean_html(issue.get("body", ""))
                    if not question or not body or len(body) < 50: continue
                    
                    comments_url = issue.get("comments_url")
                    comments_response = make_request(comments_url, headers=headers)
                    answer_parts = []
                    if comments_response:
                        comments_data = comments_response.json()
                        for comment in comments_data:
                            comment_body = clean_html(comment.get("body", ""))
                            if comment_body: answer_parts.append(comment_body)
                    
                    if answer_parts:
                        answer = "\n\n".join(answer_parts)
                        save_qa_pair(
                            question=f"{question}: {body}", answer=answer, category="devsecops",
                            subcategory="github", source=f"github_{repo_name}"
                        )
            time.sleep(random.uniform(1, 3))
    logger.info("Collecte des données GitHub terminée.")

def collect_huggingface_data(queries):
    logger.info("Début de la collecte des données Hugging Face...")
    base_url = "https://huggingface.co/api"
    headers = {"Accept": "application/json"}
    if config.USE_API_KEYS:
        token = os.getenv('HUGGINGFACE_API_TOKEN')
        headers["Authorization"] = f"Bearer {token}"
    
    search_queries = queries.split('\n') if queries else ["security", "devsecops"]
    for query in search_queries:
        logger.info(f"Recherche de datasets pour: {query}")
        search_url = f"{base_url}/datasets"
        params = {"search": query, "limit": 10}
        response = make_request(search_url, headers=headers, params=params)
        if not response: continue
        
        data = response.json()
        for dataset in data:
            dataset_id = dataset["id"].replace("/", "_")
            logger.info(f"Traitement du dataset: {dataset['id']}")
            dataset_url = f"{base_url}/datasets/{dataset['id']}"
            dataset_response = make_request(dataset_url, headers=headers)
            
            if dataset_response:
                dataset_data = dataset_response.json()
                description = clean_html(dataset_data.get("description", ""))
                if not description or len(description) < 100: continue
                tags = dataset_data.get("tags", [])
                tags_text = ", ".join(tags) if tags else "No tags"
                answer = f"Dataset: {dataset_data.get('id', '')}\nDownloads: {dataset_data.get('downloads', 0)}\nTags: {tags_text}\n\n{description}"
                
                save_qa_pair(
                    question=f"What is the {dataset_data.get('id', '')} dataset about?", answer=answer,
                    category="security", subcategory="dataset", source=f"huggingface_{dataset_id}", tags=tags
                )
            time.sleep(random.uniform(1, 3))
    logger.info("Collecte des données Hugging Face terminée.")

def collect_nvd_data():
    logger.info("Début de la collecte des données NVD...")
    base_url = "https://services.nvd.nist.gov/rest/json/cves/2.0"
    headers = {"Accept": "application/json"}
    if config.USE_API_KEYS:
        key = os.getenv('NVD_API_KEY')
        headers["apiKey"] = key
    
    params = {"resultsPerPage": 50}
    response = make_request(base_url, headers=headers, params=params)
    if not response:
        logger.warning("Impossible de récupérer les données du NVD.")
        return
    
    data = response.json()
    vulnerabilities = data.get("vulnerabilities", [])
    logger.info(f"Traitement de {len(vulnerabilities)} vulnérabilités...")
    
    for vuln in vulnerabilities:
        cve_data = vuln.get("cve", {})
        cve_id = cve_data.get("id", "")
        descriptions = cve_data.get("descriptions", [])
        description = next((desc.get("value", "") for desc in descriptions if desc.get("lang") == "en"), "")
        if not description or len(description) < 50: continue
        
        cvss_v3 = cve_data.get("metrics", {}).get("cvssMetricV31", [{}])[0].get("cvssData", {})
        severity = cvss_v3.get("baseSeverity", "UNKNOWN")
        score = cvss_v3.get("baseScore", 0)
        references = [ref.get("url", "") for ref in cve_data.get("references", [])]
        
        answer = f"CVE ID: {cve_id}\nSeverity: {severity}\nCVSS Score: {score}\nReferences: {', '.join(references[:5])}\n\nDescription: {description}"
        
        save_qa_pair(
            question=f"What is the vulnerability {cve_id}?", answer=answer,
            category="security", subcategory="vulnerability", source=f"nvd_{cve_id}"
        )
    logger.info("Collecte des données NVD terminée.")

def collect_stack_exchange_data(queries):
    logger.info("Début de la collecte des données Stack Exchange...")
    base_url = "https://api.stackexchange.com/2.3"
    params_base = {"pagesize": 10, "order": "desc", "sort": "votes", "filter": "withbody"}
    if config.USE_API_KEYS:
        key = os.getenv('STACK_EXCHANGE_API_KEY')
        params_base["key"] = key
    
    sites = [
        {"site": "security", "category": "security", "subcategory": "security"},
        {"site": "devops", "category": "devsecops", "subcategory": "devops"}
    ]
    
    tags_by_site = {
        "security": ["security", "vulnerability"],
        "devops": ["devops", "ci-cd"]
    }
    
    for site_config in sites:
        site = site_config["site"]
        category = site_config["category"]
        subcategory = site_config["subcategory"]
        logger.info(f"Collecte des données du site: {site}")
        tags = tags_by_site.get(site, []) + (queries.split('\n') if queries else [])
        
        for tag in list(set(tags)):
            logger.info(f"Recherche de questions avec le tag: {tag}")
            questions_url = f"{base_url}/questions"
            params = {**params_base, "site": site, "tagged": tag}
            
            response = make_request(questions_url, params=params)
            if not response: continue
            
            questions_data = response.json()
            for question in questions_data.get("items", []):
                question_id = question.get("question_id")
                title = question.get("title", "")
                body = clean_html(question.get("body", ""))
                if not body or len(body) < 50: continue
                
                answers_url = f"{base_url}/questions/{question_id}/answers"
                answers_params = {**params_base, "site": site}
                answers_response = make_request(answers_url, params=answers_params)
                answer_body = ""
                if answers_response and answers_response.json().get("items"):
                    answer_body = clean_html(answers_response.json()["items"][0].get("body", ""))
                
                if answer_body:
                    save_qa_pair(
                        question=title, answer=answer_body, category=category,
                        subcategory=subcategory, source=f"{site}_{question_id}", tags=question.get("tags", [])
                    )
                time.sleep(random.uniform(1, 3))
    logger.info("Collecte des données Stack Exchange terminée.")

def run_data_collection(sources, queries):
    st.session_state.bot_status = "En cours d'exécution"
    st.session_state.logs = []
    
    check_api_keys()
    
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    enabled_sources = [s for s, enabled in sources.items() if enabled]
    total_sources = len(enabled_sources)
    completed_sources = 0
    
    for source_name in enabled_sources:
        status_text.text(f"Collecte des données de {source_name}...")
        try:
            if source_name == "Kaggle":
                collect_kaggle_data(queries.get("Kaggle", ""))
            elif source_name == "GitHub":
                collect_github_data(queries.get("GitHub", ""))
            elif source_name == "Hugging Face":
                collect_huggingface_data(queries.get("Hugging Face", ""))
            elif source_name == "NVD":
                collect_nvd_data()
            elif source_name == "Stack Exchange":
                collect_stack_exchange_data(queries.get("Stack Exchange", ""))
        except Exception as e:
            logger.error(f"Erreur fatale lors de la collecte de {source_name}: {str(e)}")
        
        completed_sources += 1
        progress_bar.progress(completed_sources / total_sources)
    
    st.session_state.bot_status = "Arrêté"
    st.info("Collecte des données terminée!")
    progress_bar.empty()
    status_text.empty()

def main():
    st.title("DevSecOps Data Bot")
    st.markdown("""
        Ce bot est conçu pour collecter des données de diverses sources (GitHub, Kaggle, Hugging Face, NVD, Stack Exchange)
        afin de construire un jeu de données de questions/réponses DevSecOps.
    """)

    tabs = st.tabs(["Bot", "Statistiques & Données", "Configuration"])
    
    with tabs[0]:
        st.header("État du bot")
        col1, col2, col3 = st.columns(3)
        with col1:
            st.metric("Statut", st.session_state.bot_status)
        with col2:
            st.metric("Paires Q/R", st.session_state.total_qa_pairs)
        with col3:
            st.metric("Statut du serveur LLM", st.session_state.server_status)
        
        st.markdown("---")
        
        st.header("Lancer la collecte")
        
        st.subheader("Sources de données")
        sources_columns = st.columns(5)
        sources = {
            "GitHub": sources_columns[0].checkbox("GitHub", value=True),
            "Kaggle": sources_columns[1].checkbox("Kaggle", value=True),
            "Hugging Face": sources_columns[2].checkbox("Hugging Face", value=True),
            "NVD": sources_columns[3].checkbox("NVD", value=True),
            "Stack Exchange": sources_columns[4].checkbox("Stack Exchange", value=True),
        }
        
        st.subheader("Requêtes de recherche")
        queries = {}
        queries["GitHub"] = st.text_area("Requêtes GitHub (une par ligne)", "topic:devsecops\ntopic:security\nvulnerability")
        queries["Kaggle"] = st.text_area("Requêtes Kaggle (une par ligne)", "cybersecurity\nvulnerability dataset\npenetration testing")
        queries["Hugging Face"] = st.text_area("Requêtes Hugging Face (une par ligne)", "security dataset\nvulnerability\nlanguage model security")
        queries["Stack Exchange"] = st.text_area("Tags Stack Exchange (un par ligne)", "devsecops\nsecurity\nvulnerability")
        
        st.markdown("---")
        
        if st.session_state.bot_status == "Arrêté":
            if st.button("Lancer la collecte", use_container_width=True, type="primary"):
                st.session_state.logs = []
                st.session_state.qa_data = []
                st.session_state.total_qa_pairs = 0
                run_data_collection(sources, queries)
        else:
            st.warning("La collecte est en cours. Veuillez attendre qu'elle se termine.")
            if st.button("Forcer l'arrêt", use_container_width=True, type="secondary"):
                st.session_state.bot_status = "Arrêté"
                st.info("La collecte a été arrêtée manuellement.")
        
        st.markdown("---")
        st.subheader("Logs d'exécution")
        log_container = st.container(border=True)
        with log_container:
            for log in st.session_state.logs:
                st.text(log)
    
    with tabs[1]:
        st.header("Statistiques")
        if st.session_state.qa_data:
            df = pd.DataFrame(st.session_state.qa_data)
            
            st.subheader("Aperçu des données")
            st.dataframe(df, use_container_width=True)
            
            st.subheader("Répartition par source")
            source_counts = df['source'].apply(lambda x: x.split('_')[0]).value_counts().reset_index()
            source_counts.columns = ['Source', 'Nombre']
            fig_source = px.bar(source_counts, x='Source', y='Nombre', title="Nombre de paires Q/R par source")
            st.plotly_chart(fig_source, use_container_width=True)
            
            st.subheader("Répartition par catégorie")
            category_counts = df['category'].value_counts().reset_index()
            category_counts.columns = ['Catégorie', 'Nombre']
            fig_cat = px.pie(category_counts, names='Catégorie', values='Nombre', title="Répartition par catégorie")
            st.plotly_chart(fig_cat, use_container_width=True)
            
            st.subheader("Téléchargement des données")
            col1, col2 = st.columns(2)
            with col1:
                json_data = json.dumps(st.session_state.qa_data, indent=2)
                st.download_button(
                    label="Télécharger JSON",
                    data=json_data,
                    file_name=f"devsecops_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                    mime="application/json",
                    use_container_width=True
                )
            with col2:
                csv_data = df.to_csv(index=False)
                st.download_button(
                    label="Télécharger CSV",
                    data=csv_data,
                    file_name=f"devsecops_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv",
                    mime="text/csv",
                    use_container_width=True
                )
        else:
            st.info("Aucune donnée à afficher. Lancez d'abord la collecte de données.")
    
    with tabs[2]:
        st.header("Configuration Avancée")
        st.subheader("Paramètres du serveur LLM")
        
        llm_col1, llm_col2 = st.columns(2)
        with llm_col1:
            if st.button("Démarrer le serveur LLM", type="primary", use_container_width=True):
                start_llm_server()
            if st.button("Vérifier le statut du serveur", use_container_width=True):
                check_server_status()
                st.rerun()
        with llm_col2:
            if st.button("Arrêter le serveur LLM", type="secondary", use_container_width=True):
                stop_llm_server()
        
        st.markdown("---")
        
        st.subheader("Paramètres d'enrichissement IA")
        st.session_state.enable_enrichment = st.checkbox("Activer l'enrichissement IA", value=st.session_state.enable_enrichment, help="Utilise le LLM pour améliorer les paires Q/R.")
        st.session_state.min_relevance = st.slider("Score de pertinence minimum", 0, 100, st.session_state.min_relevance, help="Les contenus en dessous de ce score ne seront pas traités.")
        st.session_state.temperature = st.slider("Température de l'IA", 0.0, 1.0, st.session_state.temperature, help="Contrôle la créativité de l'IA. 0.0 = plus déterministe, 1.0 = plus créatif.")
        st.session_state.n_predict = st.slider("Nombre de tokens", 128, 1024, st.session_state.n_predict, help="Nombre maximum de tokens à générer par l'IA.")

if __name__ == "__main__":
    create_initial_setup()
    main()