from xml.parsers.expat import model

import requests
import json
import os
import time
import tqdm
import re
from dotenv import load_dotenv
import chromadb
from sentence_transformers import SentenceTransformer
import uuid
# Load variables from .env file
load_dotenv()

# --- Configuration ---
DUST_API_KEY = os.getenv("DUST_API_KEY")  
WORKSPACE_ID = os.getenv("DUST_WORKSPACE_ID")
AGENT_ID = "dust"  # ou l'identifiant de ton agent spécifique

BASE_URL = "https://dust.tt/api/v1"


model = SentenceTransformer("all-MiniLM-L6-v2")
HEADERS = {
    "Authorization": f"Bearer {DUST_API_KEY}",
    "Content-Type": "application/json",
}

client = chromadb.PersistentClient(path="./chroma_db")
COLLECTION_NAME = "abap_rag"
collection = client.get_or_create_collection(
    name=COLLECTION_NAME
)



def create_conversation(message: str) -> dict:
    """Crée une nouvelle conversation et envoie un premier message."""
    url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations"

    payload = {
        "message": {
            "content": message,
            "mentions": [{"configurationId": AGENT_ID}],
            "context": {
                "timezone": "Europe/Paris",
                "username": "python_user",
                "fullName": "Python User",
                "email": "python@example.com",
                "profilePictureUrl": None,
            },
        },
        "visibility": "unlisted",
        "title": None,
    }

    response = requests.post(url, headers=HEADERS, json=payload)
    response.raise_for_status()
    return response.json()


def get_agent_message_events(conversation_id: str, message_id: str):
    """Récupère la réponse de l'agent en streaming (SSE)."""
    url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations/{conversation_id}/messages/{message_id}/events"

    with requests.get(url, headers=HEADERS, stream=True) as response:
        response.raise_for_status()
        full_text = ""

        for line in response.iter_lines():
            if line:
                decoded = line.decode("utf-8")
                if decoded.startswith("data:"):
                    data_str = decoded[len("data:"):].strip()
                    try:
                        event = json.loads(data_str)
                        event_type = event.get("data").get("type")

                        if event_type == "generation_tokens":
                            token = event.get("data").get("text", "")
                            print(token, end="", flush=True)
                            full_text += token

                        elif event_type == "agent_message_success":
                            print()  # Saut de ligne final
                            break

                        elif event_type == "agent_error":
                            print(f"\n[Erreur agent] {event}")
                            break

                    except json.JSONDecodeError:
                        pass

        return full_text


def send_followup_message(conversation_id: str, message: str) -> dict:
    """Envoie un message de suivi dans une conversation existante."""
    url = f"{BASE_URL}/w/{WORKSPACE_ID}/assistant/conversations/{conversation_id}/messages"

    payload = {
        "content": message,
        "mentions": [{"configurationId": AGENT_ID}],
        "context": {
            "timezone": "Europe/Paris",
            "username": "python_user",
            "fullName": "Python User",
            "email": "python@example.com",
            "profilePictureUrl": None,
        },
    }

    response = requests.post(url, headers=HEADERS, json=payload)
    response.raise_for_status()
    return response.json()


def chat(message: str):
    """Point d'entrée principal : crée une conversation et affiche la réponse."""
    print(f"\n[Vous] {message}")
    print("[Agent] ", end="")

    # 1. Créer la conversation
    result = create_conversation(message)
    conversation = result.get("conversation", {})
    conversation_id = conversation.get("sId")

    # 2. Trouver l'ID du message agent dans la conversation
    agent_message = None
    for msg in conversation.get("content", []):
        for m in msg:
            if m.get("type") == "agent_message":
                agent_message = m
                break

    if not agent_message:
        print("Aucun message agent trouvé.")
        return None

    agent_message_id = agent_message.get("sId")

    # 3. Streamer la réponse
    response_text = get_agent_message_events(conversation_id, agent_message_id)

    return {"conversation_id": conversation_id, "response": response_text}

# =========================
# INDEX FILES
# =========================

def index_folder(folder_path):

    all_files = []

    for root, dirs, files in os.walk(folder_path):
        for file in files:

            if file.endswith(".abap"):

                all_files.append(
                    os.path.join(root, file)
                )

    print(f"{len(all_files)} fichiers trouvés")

    for filepath in tqdm.tqdm(all_files):

        try:

            with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
                content = f.read()

            chunks = split_abap_code(content)

            for chunk in chunks:

                embedding = model.encode(chunk).tolist()

                collection.add(
                    ids=[str(uuid.uuid4())],
                    documents=[chunk],
                    embeddings=[embedding],
                    metadatas=[{
                        "source": filepath,
                        "filename": os.path.basename(filepath)
                    }]
                )

        except Exception as e:
            print(f"Erreur indexation {filepath}: {e}")

    print("Indexation terminée") 
# =========================
# DOCUMENT FILE WITH RAG
# =========================

def document_file(filepath):

    with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
        content = f.read()

    filename = os.path.basename(filepath)

    # recherche RAG
    rag_context = retrieve_context(content, top_k=5)

    prompt = f"""
Tu es un expert SAP ABAP.

Tu dois documenter le fichier suivant.

FICHIER ANALYSÉ:
{filename}

CHEMIN:
{filepath}

CODE PRINCIPAL:
{content}

CONTEXTE RAG:
{rag_context}

Ton résultat doit être en markdown avec plusieurs sections :

# Algorithmie
Explique le fonctionnement général du programme.

# Objets
Explique les objets utilisés et leur rôle.

# Dépendances
Explique les dépendances importantes.

# Flux de traitement
Décris les principales étapes d'exécution.

# Points d'attention
Liste les risques techniques ou métiers.
"""

    return chat(prompt)
# =========================
# ABAP CHUNKER
# =========================

def split_abap_code(content):
    """
    Découpe :
    - METHOD
    - FORM
    - FUNCTION
    """

    patterns = [
        r"METHOD\s+.*?ENDMETHOD\.",
        r"FORM\s+.*?ENDFORM\.",
        r"FUNCTION\s+.*?ENDFUNCTION\."
    ]

    chunks = []

    for pattern in patterns:
        matches = re.findall(
            pattern,
            content,
            re.IGNORECASE | re.DOTALL
        )

        chunks.extend(matches)

    # fallback si aucun chunk
    if not chunks:
        chunks = [content]

    return chunks
# =========================
# RAG SEARCH
# =========================

def retrieve_context(query, top_k=5):

    query_embedding = model.encode(query).tolist()

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )

    documents = results["documents"][0]
    metadatas = results["metadatas"][0]

    context_parts = []

    for doc, meta in zip(documents, metadatas):

        source = meta.get("source", "unknown")

        context_parts.append(f"""
SOURCE FILE:
{source}

CODE:
{doc}
""")

    return "\n\n".join(context_parts)
# --- Programme principal ---
if __name__ == "__main__":
    folder_or_file = input("Entrez le chemin d'un dossier ou d'un fichier à analyser : ")
    outputdir = input("Entrez le chemin du dossier de sortie pour les documentations générées : ")
    if os.path.isdir(folder_or_file):
        all_files = []
    
        for root, dirs, files in os.walk(folder_or_file):
            for file in files:
                if file.endswith(".abap"):
                    all_files.append(
                        os.path.join(root, file)
                    )
        print("INDEXATION RAG...")
        index_folder(folder_or_file)

        print("GÉNÉRATION DOCUMENTATION...")

        for root, dirs, files in os.walk(folder_or_file):

            for file in files:

                if file.endswith(".abap"):

                    filepath = os.path.join(root, file)

                    print(f"\nDocumentation de {filepath}")

                    result = document_file(filepath)

                    if result:

                        documentation = result['response']

                        output_name = os.path.splitext(file)[0]

                        with open(
                            f"{outputdir}/documentation_{output_name}.md",
                            "w",
                            encoding="utf-8"
                        ) as f:

                            f.write(documentation)

                        print(f"Documentation générée : documentation_{output_name}.md")
    elif os.path.isfile(folder_or_file):
        content = open(folder_or_file, "r").read()
        result = chat(f'''peux-tu m'expliquer le fonctionnement du fichier {content} ton résultat devra être en markdown avec plusieurs sections :
                    une section algorithmie, où tu explique le fonctionnement général du programme 
                    une section objets, où tu explique les différents objets utilisés dans le programme et leur rôle.''' )

        if result:
            print(f"\n[réponse: {result}]")
            print(f"\n[Conversation ID: {result['conversation_id']}]")
            documentation = result['response']
            file_name = os.path.basename(folder_or_file).split('.')[0]
            with open(f"documentation_{file_name}.md", "w", encoding="utf-8") as f:
                f.write(documentation)
    else:
        print("Le chemin spécifié n'est ni un fichier ni un dossier valide.")