import os
import json

from langchain.document_loaders import PyPDFLoader
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain  # Model and API configuration

# /!\ Trigger Warning: this is a first commit the code is not supposed to work in the library at the moment
# updates in progress

# keys and endpoints
MODEL_NAME = "gpt-4o-mini"  # OpenRouter GPT-4o model
OPENROUTER_API_BASE = "https://openrouter.ai/api/v1"
OPENROUTER_API_KEY = "" # your key is needed my friend!
temperature=0 # Initialize LLM once at import time

# checking LLM Client initialisation
try:
    llm = ChatOpenAI(
        model_name=MODEL_NAME,
        temperature=temperature,
        openai_api_base=OPENROUTER_API_BASE,
        openai_api_key=OPENROUTER_API_KEY
    )
    print("LLM initialized successfully.")
except Exception as e:
    llm = None
    print(f"Failed to initialize LLM: {e}")

# Smoke-test chat connectivity via OpenRouter 
def _run_smoke_test():
    """Quick check to verify chat connectivity with OpenRouter LLM."""
    from langchain.schema import HumanMessage
    try:
        test_prompt = "Hi there! Can you confirm you're alive?"
        resp = llm([HumanMessage(content=test_prompt)])
        print("🔹 Chat OK →", resp.content[:80], "…")
    except Exception as e:
        raise RuntimeError(f"⚠️ Chat smoke-test failed: {e}")

# Run smoke-test
_run_smoke_test()

# New: Layout extraction function
def extract_layout(pdf_path: str) -> str:
    """
    Extract the document's structural layout (sections and subsections with page numbers)
    using a single LLM call with the full context.

    Args:
        pdf_path (str): Path to the PDF file.

    Returns:
        str: Document layout description generated by the LLM.
    """
    if not llm:
        raise RuntimeError("LLM is not initialized. Please check your API key.")

    # Load full text of the PDF
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()
    full_text = "\n".join([doc.page_content for doc in documents])

    # Build prompt for single-shot layout extraction
    prompt = (
        "Extrait la structure du document en listant les sections et sous-sections "
        "avec leurs numéros de page. Présente le résultat sous forme de liste hiérarchique. "
        "Sois précis et n'invente pas d'information :\n" + full_text
    )
    # Call LLM once
    response = llm([HumanMessage(content=prompt)])
    return response.content.strip()


def analyze_pdf_unified(pdf_path: str, layout_str: str) -> str:
    """
    Perform all types of analysis on a PDF in a single LLM call using a unified prompt.

    This includes:
      1. Structured summary (by domain topic)
      2. Hierarchical summary (by layout depth)
      3. Entity extraction + Mermaid graph
      4. Concept extraction + Mermaid graph

    Args:
        pdf_path (str): Path to the PDF file.
        layout_str (str): Pre-extracted layout (as plain text)

    Returns:
        str: Unified analysis response
    """
    if not llm:
        raise RuntimeError("LLM is not initialized. Please check your API key.")

    # Load document content
    loader = PyPDFLoader(pdf_path)
    docs = loader.load()
    full_text = "\n".join([doc.page_content for doc in docs])

    # Construct unified prompt
    prompt = f"""
    Vous allez recevoir la structure du document et son contenu intégral. Votre mission est d'en produire une analyse complète, en une seule réponse, structurée selon les instructions suivantes.

    ---

    ## Entrées

    **Structure :**
    {layout_str}

    **Contenu complet :**
    {full_text}

    ---

    ## Sortie attendue

    ### 1. Résumé structuré (par thématique fonctionnelle)
    Pour chacune des rubriques suivantes :
      1. Contexte et objectifs  
      2. Analyse opérationnelle  
      3. Analyse système  
      4. Architecture logique  
      5. Architecture physique  
      6. Sécurité et perspectives  
    - Résumez en 1 phrase l'objectif global de chaque rubrique  
    - Puis ajoutez 1 phrase par paragraphe important  
    - Ajoutez tout autre point saillant si pertinent

    ### 2. Résumé hiérarchique précis (par structure)
    Respectez précisément chaque niveau indiqué dans la structure (ex : 1, 1.1, 1.1.1...) :
    - Pour chaque niveau :
      - Résumé global : 1 phrase
      - Résumé détaillé : 1 phrase par paragraphe

    ### 3. Extraction hiérarchique des entités et tags
    Pour chaque niveau (section, sous-section, sous-sous-section exactement comme indiqué dans la structure) :
    - Listez les entités nommées clairement identifiées (personnes, organisations, lieux, technologies)
    - Proposez des tags pertinents spécifiques à chaque niveau

    Générez ensuite un diagramme Mermaid hiérarchisé clair, par exemple :

    ```mermaid
    graph TD
      1["1. Titre Section"] --> 1_ENT1["Organisation: ACTIA"]
      1 --> 1_TAG1["Tag: Technologie"]
      1 --> 1.1["1.1 Titre Sous-section"]
      1.1 --> 1.1_ENT1["Personne: Jean Dupont"]
    ```

    ### 4. Extraction hiérarchique des concepts
        Pour chaque niveau indiqué dans la structure :
        - Listez les concepts-clés (méthodologies, composants, technologies, principes techniques) qui se trouvent dans les paragraphes

    Générez un diagramme Mermaid hiérarchisé correspondant :

    ```mermaid
    graph TD
          1["1. Titre Section"] --> 1_CPT1["Concept: Machine Learning"]
          1 --> 1.1["1.1 Titre Sous-section"]
      1.1 --> 1.1_CPT1["Concept: GPT-4"]
    ```

    ---

    Répondez dans l’ordre exact ci-dessus. Soyez clair, synthétique, précis et strictement fidèle au texte, sans inventer d'information.
    """

    # Call LLM once
    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content.strip()

# Function to find relevant sections based on multiple criteria including Mermaid graph
def find_relevant_sections_detailed(analysis_result: str, layout_str: str, user_query: str) -> str:
    prompt = f"""
    À partir de l'analyse complète suivante du document, incluant les sections, les entités, les tags, et les concepts (présentés en format Mermaid), identifie précisément les sections pertinentes pour répondre à la question suivante:

    Analyse complète avec graph Mermaid :
    {analysis_result}

    Structure du document :
    {layout_str}

    Question :
    {user_query}

    Utilise explicitement la hiérarchie (chapitres, sections, jusqu'aux sous-sections), les résumés ainsi que ce qui a été trouvé en MERMAID: les noms, les résumés, les tags, entités et concepts pour identifier et justifier les sections pertinentes. Réponds avec les numéros et titres exacts des sections pertinentes ainsi qu'une courte justification explicite de leur pertinence basée sur les tags, entités, concepts et résumés fournis.
    
    produit un mermaid des liens trouvés entre les résumés/chapitres, sections, sous sections, noms, concepts tags, entités et concepts afin d'exprimer ton raisonement
    
    à la fin liste les sections et sous-sections utilisées, il est important d'avoir les 2.
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content.strip()

# Extract raw text from sections using LLM call
def extract_raw_text_sections_llm(pdf_path: str, sections: str) -> str:
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()
    full_text = "\n".join([doc.page_content for doc in documents])

    prompt = f"""
    Extrait précisément le texte brut des sections et sous-sections suivantes à partir du document fourni :

    Sections et sous-sections :
    {sections}

    Texte complet :
    {full_text}

    Retourne uniquement le texte brut exact des sections demandées, sans ajouter ni modifier d'information.
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content.strip()


# Generate questions based on sections and Mermaid graph
def generate_questions_from_sections(final_sections: str, analysis_result: str, user_query: str, categories_filepath: str) -> str:
    
    with open(categories_filepath, 'r', encoding='utf-8') as file:
        categories_definition = file.read()
        
    prompt = f"""
    À partir du texte extrait des sections suivantes :
    {final_sections}

    et de l'analyse complète (incluant le graphe Mermaid) :
    {analysis_result}

    Catégories de questions :
    {categories_definition}

    Génère des questions pertinentes pour chaque section et sous-section en lien direct avec la requête utilisateur suivante :
    {user_query}

    Ensuite, utilise les liens indiqués dans le graphe Mermaid pour créer des questions additionnelles reliant ces sections, concepts, entités entre elles.

    Structure clairement chaque série de questions par section, entités, concepts et indique explicitement les liens utilisés pour générer les questions additionnelles.
    Pour chaque question précise les concepts, entités utilisées.
    
    enfin de toute ces questions essais d'en générer des plus globales toujours en liant sections, entités, concepts.
    
    n'oublie pas de classifier chaque question selon les catégories fournies dans le fichier
    
    remember those question must be related to concepts and entities of the user input
    
    n'oublie pas de générer aussi les réponses à chaque fois!
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content.strip()

# Refine questions based on user query, concepts, and categories
def refine_questions(user_query: str, questions_generated: str, categories_filepath: str) -> str:
    with open(categories_filepath, 'r', encoding='utf-8') as file:
        categories_definition = file.read()

    prompt = f"""
    Question utilisateur initiale :
    {user_query}

    Questions générées :
    {questions_generated}

    Catégories de questions :
    {categories_definition}

    Sélectionne les 5 questions les plus pertinentes par rapport à la question utilisateur initiale, en tenant compte de la proximité sémantique, des concepts clés abordés, les entités et en assurant une diversité des catégories représentées.
    Les questions doivent avoir une proximité des concepts et entités avec les questions générées.
    Retourne ces 5 questions raffinées clairement classifiées selon les catégories fournies ainsi que les réponses à ne pas oublier.
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return response.content.strip()


# Ton chemin vers le fichier PDF
pdf_path = "data/AIDA Architecture synthesis V4.5.pdf"

# extraction layout
extracted_layout =  extract_layout(pdf_path)

# Appel à la fonction unifiée
result = analyze_pdf_unified(pdf_path, extracted_layout)


categories_filepath = "data/overview_question_categories.txt"

questions = [
    "What are the main components of the system described in the document?",
    "High level system capabilities and activities?",
    "What is the definition of the acronyms in the document?",
    "What are the three main segments of the system as described in the document?",
    "What are the primary missions of the system as outlined in the document?",
    "What are the identified failure scenarios for the system as described in the document?",
    "What are the high-level functions of the system as described in the document, and how do they contribute to achieving the system's overall objectives?",
    "What are the high-level functions for the logical layer of the system as described in the document, and how do they contribute to the overall system architecture?",
    "What are the main components of the system described in the document?",
    "Can you list the subsystems of the propulsion component?",
    "Which subsytem is the most complex?",
    "What are the key functions and safety considerations of the propulsion system in the AIDA?",
    "What are the key safety considerations and mitigation strategies for the AIDA system as described in the document?",
    "What is the interaction between components 'The remote control' and 'Flight control and monitoring system'?",
    "What are the detailed specifications for the AIDA system's propulsion units that allow it to achieve supersonic speeds?",
    "How does the AIDA system's propulsion system enable it to operate underwater for marine inspections?",
    "What are the specific battery technologies used in the AIDA system's power supply, and how do they compare in terms of performance and cost?",
    "Can you provide the exact dimensions and weight of the AIDA drone as specified in the document?",
    "Can you provide a detailed analysis of the AIDA system's flight data over the past year, including trends in performance metrics and failure rates?",
    "What is the comprehensive impact assessment of integrating the AIDA system into existing airport infrastructure, considering factors such as operational efficiency, safety improvements, and cost savings?"
]

refined_questions_list = []

for user_query in questions:
    print(user_query)
    relevant_sections_detailed = find_relevant_sections_detailed(result, extracted_layout, user_query)
    final_extracted_sections = extract_raw_text_sections_llm(pdf_path, relevant_sections_detailed)
    questions_generated = generate_questions_from_sections(final_extracted_sections, result, user_query, categories_filepath)
    refined_questions_list.append(refine_questions(user_query, questions_generated, categories_filepath))