import os
import logging
from typing import List
from pydantic import BaseModel
from fastapi import FastAPI, HTTPException
import rdflib
from rdflib import RDF, RDFS, OWL
from sentence_transformers import SentenceTransformer
import faiss
import json
import numpy as np
from dotenv import load_dotenv
import requests

# Carica le variabili d'ambiente
load_dotenv()

# Configura il logging
logging.basicConfig(
    level=logging.INFO,  # Ridotto da DEBUG per migliorare le prestazioni
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.FileHandler("app.log"), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# Recupera la chiave API
API_KEY = os.getenv("HF_API_KEY")
if not API_KEY:
    logger.error("HF_API_KEY non impostata.")
    raise EnvironmentError("HF_API_KEY non impostata.")

# Definisci i percorsi dei file
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
RDF_FILE = os.path.join(BASE_DIR, "Ontologia.rdf")
HF_MODEL = "google/flan-t5-xxl"  # Modello ottimizzato per seguire istruzioni

MAX_CLASSES = 30
MAX_PROPERTIES = 30

# Percorsi dei file generati
DOCUMENTS_FILE = os.path.join(BASE_DIR, "data", "documents.json")
FAISS_INDEX_FILE = os.path.join(BASE_DIR, "data", "faiss.index")

def create_data_directory():
    """Crea la directory 'data/' se non esiste."""
    os.makedirs(os.path.join(BASE_DIR, "data"), exist_ok=True)
    logger.info("Directory 'data/' creata o già esistente.")

def extract_ontology(rdf_file: str, output_file: str):
    """
    Estrae classi, proprietà ed entità dall'ontologia RDF e le salva in un file JSON come un unico documento.
    """
    logger.info(f"Inizio estrazione dell'ontologia da {rdf_file}.")
    g = rdflib.Graph()
    try:
        g.parse(rdf_file, format="xml")
        logger.info(f"Parsing RDF di {rdf_file} riuscito.")
    except Exception as e:
        logger.error(f"Errore nel parsing RDF: {e}")
        raise e

    # Estrai Classi
    classes = []
    for cls in g.subjects(RDF.type, OWL.Class):
        label = g.value(cls, RDFS.label, default=str(cls))
        description = g.value(cls, RDFS.comment, default="No description.")
        classes.append({"class": str(cls), "label": str(label), "description": str(description)})

    for cls in g.subjects(RDF.type, RDFS.Class):
        label = g.value(cls, RDFS.label, default=str(cls))
        description = g.value(cls, RDFS.comment, default="No description.")
        classes.append({"class": str(cls), "label": str(label), "description": str(description)})

    # Estrai Proprietà
    properties = []
    for prop in g.subjects(RDF.type, OWL.ObjectProperty):
        label = g.value(prop, RDFS.label, default=str(prop))
        description = g.value(prop, RDFS.comment, default="No description.")
        properties.append({"property": str(prop), "label": str(label), "description": str(description)})

    for prop in g.subjects(RDF.type, OWL.DatatypeProperty):
        label = g.value(prop, RDFS.label, default=str(prop))
        description = g.value(prop, RDFS.comment, default="No description.")
        properties.append({"property": str(prop), "label": str(label), "description": str(description)})

    for prop in g.subjects(RDF.type, RDF.Property):
        label = g.value(prop, RDFS.label, default=str(prop))
        description = g.value(prop, RDFS.comment, default="No description.")
        properties.append({"property": str(prop), "label": str(label), "description": str(description)})

    # Estrai Entità (NamedIndividuals)
    entities = []
    for entity in g.subjects(RDF.type, OWL.NamedIndividual):
        label = g.value(entity, RDFS.label, default=str(entity))
        description = g.value(entity, RDFS.comment, default="No description.")
        # Estrai le proprietà dell'entità
        entity_properties = {}
        for predicate, obj in g.predicate_objects(entity):
            if predicate not in [RDFS.label, RDFS.comment]:
                entity_properties[str(predicate)] = str(obj)
        entities.append({
            "entity": str(entity),
            "label": str(label),
            "description": str(description),
            "properties": entity_properties
        })

    # Crea un unico documento
    ontology_summary = {
        "title": "Ontologia Museo",
        "classes": classes[:MAX_CLASSES],
        "properties": properties[:MAX_PROPERTIES],
        "entities": entities,  # Aggiungi le entità
        "full_ontology": g.serialize(format="xml").decode('utf-8') if isinstance(g.serialize(format="xml"), bytes) else g.serialize(format="xml")  # Decodifica rimossa
    }

    # Salva il documento in JSON
    try:
        with open(output_file, "w", encoding="utf-8") as f:
            json.dump(ontology_summary, f, ensure_ascii=False, indent=2)
        logger.info(f"Ontologia estratta e salvata in {output_file}")
    except Exception as e:
        logger.error(f"Errore nel salvataggio di {output_file}: {e}")
        raise e

def create_faiss_index(documents_file: str, index_file: str, embedding_model: str = 'all-MiniLM-L6-v2'):
    """
    Crea un indice FAISS a partire dal documento estratto.
    """
    logger.info(f"Inizio creazione dell'indice FAISS da {documents_file}.")
    try:
        # Carica il documento
        with open(documents_file, "r", encoding="utf-8") as f:
            document = json.load(f)
        logger.info(f"Documento caricato da {documents_file}.")

        # Genera embedding
        model = SentenceTransformer(embedding_model)
        # Concatenazione delle classi, proprietà e entità per l'embedding
        texts = [f"Classe: {cls['label']}. Descrizione: {cls['description']}" for cls in document['classes']]
        texts += [f"Proprietà: {prop['label']}. Descrizione: {prop['description']}" for prop in document['properties']]
        texts += [f"Entità: {entity['label']}. Descrizione: {entity['description']}. Proprietà: {entity['properties']}" for entity in document.get('entities', [])]
        embeddings = model.encode(texts, convert_to_numpy=True)
        logger.info("Embedding generati con SentenceTransformer.")

        # Crea l'indice FAISS
        dimension = embeddings.shape[1]
        index = faiss.IndexFlatL2(dimension)
        index.add(embeddings)
        logger.info(f"Indice FAISS creato con dimensione: {dimension}.")

        # Salva l'indice
        faiss.write_index(index, index_file)
        logger.info(f"Indice FAISS salvato in {index_file}.")
    except Exception as e:
        logger.error(f"Errore nella creazione dell'indice FAISS: {e}")
        raise e

def prepare_retrieval():
    """Prepara i file necessari per l'approccio RAG."""
    logger.info("Inizio preparazione per il retrieval.")
    create_data_directory()

    # Verifica se Ontologia.rdf esiste
    if not os.path.exists(RDF_FILE):
        logger.error(f"File RDF non trovato: {RDF_FILE}")
        raise FileNotFoundError(f"File RDF non trovato: {RDF_FILE}")
    else:
        logger.info(f"File RDF trovato: {RDF_FILE}")

    # Verifica se documents.json esiste, altrimenti generarlo
    if not os.path.exists(DOCUMENTS_FILE):
        logger.info(f"File {DOCUMENTS_FILE} non trovato. Estrazione dell'ontologia.")
        try:
            extract_ontology(RDF_FILE, DOCUMENTS_FILE)
        except Exception as e:
            logger.error(f"Errore nell'estrazione dell'ontologia: {e}")
            raise e
    else:
        logger.info(f"File {DOCUMENTS_FILE} trovato.")

    # Verifica se faiss.index esiste, altrimenti crearlo
    if not os.path.exists(FAISS_INDEX_FILE):
        logger.info(f"File {FAISS_INDEX_FILE} non trovato. Creazione dell'indice FAISS.")
        try:
            create_faiss_index(DOCUMENTS_FILE, FAISS_INDEX_FILE)
        except Exception as e:
            logger.error(f"Errore nella creazione dell'indice FAISS: {e}")
            raise e
    else:
        logger.info(f"File {FAISS_INDEX_FILE} trovato.")

def extract_classes_and_properties(rdf_file: str) -> str:
    """
    Carica l'ontologia e crea un 'sunto' di Classi, Proprietà ed Entità
    (senza NamedIndividuals) per ridurre i token.
    """
    logger.info(f"Inizio estrazione di classi, proprietà ed entità da {rdf_file}.")
    g = rdflib.Graph()
    try:
        g.parse(rdf_file, format="xml")
        logger.info(f"Parsing RDF di {rdf_file} riuscito.")
    except Exception as e:
        logger.error(f"Errore nel parsing RDF: {e}")
        return "PARSING_ERROR"

    # Troviamo le classi
    classes_found = set()
    for s in g.subjects(RDF.type, OWL.Class):
        classes_found.add(s)
    for s in g.subjects(RDF.type, RDFS.Class):
        classes_found.add(s)
    classes_list = sorted(str(c) for c in classes_found)
    classes_list = classes_list[:MAX_CLASSES]

    # Troviamo le proprietà
    props_found = set()
    for p in g.subjects(RDF.type, OWL.ObjectProperty):
        props_found.add(p)
    for p in g.subjects(RDF.type, OWL.DatatypeProperty):
        props_found.add(p)
    for p in g.subjects(RDF.type, RDF.Property):
        props_found.add(p)
    props_list = sorted(str(x) for x in props_found)
    props_list = props_list[:MAX_PROPERTIES]

    # Troviamo le entità
    entities_found = set()
    for e in g.subjects(RDF.type, OWL.NamedIndividual):
        entities_found.add(e)
    entities_list = sorted(str(e) for e in entities_found)
    entities_list = entities_list[:MAX_CLASSES]  # Puoi impostare un limite adeguato

    txt_classes = "\n".join([f"- CLASSE: {c}" for c in classes_list])
    txt_props = "\n".join([f"- PROPRIETÀ: {p}" for p in props_list])
    txt_entities = "\n".join([f"- ENTITÀ: {e}" for e in entities_list])

    summary = f"""\
# CLASSI (max {MAX_CLASSES})
{txt_classes}
# PROPRIETÀ (max {MAX_PROPERTIES})
{txt_props}
# ENTITÀ (max {MAX_CLASSES})
{txt_entities}
"""
    logger.info("Estrazione di classi, proprietà ed entità completata.")
    return summary

def retrieve_relevant_documents(query: str, top_k: int = 5):
    """Recupera i documenti rilevanti usando FAISS."""
    logger.info(f"Recupero dei documenti rilevanti per la query: {query}")
    try:
        # Carica il documento
        with open(DOCUMENTS_FILE, "r", encoding="utf-8") as f:
            document = json.load(f)
        logger.info(f"Documento caricato da {DOCUMENTS_FILE}.")

        # Carica l'indice FAISS
        index = faiss.read_index(FAISS_INDEX_FILE)
        logger.info(f"Indice FAISS caricato da {FAISS_INDEX_FILE}.")

        # Genera embedding della query
        model = SentenceTransformer('all-MiniLM-L6-v2')
        query_embedding = model.encode([query], convert_to_numpy=True)
        logger.info("Embedding della query generati.")

        # Ricerca nell'indice
        distances, indices = index.search(query_embedding, top_k)
        logger.info(f"Ricerca FAISS completata. Risultati ottenuti: {len(indices[0])}")

        # Concatenazione delle descrizioni per la ricerca
        texts = [f"Classe: {cls['label']}. Descrizione: {cls['description']}" for cls in document['classes']]
        texts += [f"Proprietà: {prop['label']}. Descrizione: {prop['description']}" for prop in document['properties']]
        texts += [f"Entità: {entity['label']}. Descrizione: {entity['description']}. Proprietà: {entity['properties']}" for entity in document.get('entities', [])]

        # Recupera i testi rilevanti
        relevant_texts = [texts[idx] for idx in indices[0] if idx < len(texts)]
        retrieved_docs = "\n".join(relevant_texts)
        logger.info(f"Documenti rilevanti recuperati: {len(relevant_texts)}")
        return retrieved_docs
    except Exception as e:
        logger.error(f"Errore nel recupero dei documenti rilevanti: {e}")
        raise e

def create_system_message(ont_text: str, retrieved_docs: str) -> str:
    """
    Prompt di sistema robusto, con regole su query in una riga e
    informazioni recuperate tramite RAG.
    """
    return f"""\
### Istruzioni ###
Sei un assistente museale esperto in ontologie RDF. Utilizza le informazioni fornite per generare query SPARQL precise e pertinenti.

### Ontologia ###
{ont_text}
### FINE Ontologia ###

Ecco alcune informazioni rilevanti recuperate dalla base di conoscenza:
{retrieved_docs}

### Regole Stringenti ###
1) Se l'utente chiede informazioni su questa ontologia, genera SEMPRE una query SPARQL in UNA SOLA RIGA, con prefix:
   PREFIX base: <http://www.semanticweb.org/lucreziamosca/ontologies/progettoMuseo#>
2) La query SPARQL deve essere precisa e cercare esattamente le entità specificate dall'utente. Ad esempio, se l'utente chiede "Chi ha creato l'opera 'Amore e Psiche'?", la query dovrebbe cercare l'opera esattamente con quel nome.
3) Se la query produce 0 risultati o fallisce, ritenta con un secondo tentativo.
4) Se la domanda è generica (tipo 'Ciao, come stai?'), rispondi breve.
5) Se trovi risultati, la risposta finale deve essere la query SPARQL (una sola riga).
6) Se non trovi nulla, rispondi con 'Nessuna info.'
7) Non multiline. Esempio: PREFIX base: <...> SELECT ?x WHERE {{ ... }}.
Esempio:
Utente: Chi ha creato l'opera 'Amore e Psiche'?
Risposta: PREFIX base: <http://www.semanticweb.org/lucreziamosca/ontologies/progettoMuseo#> SELECT ?creatore WHERE {{ ?opera base:hasName "Amore e Psiche" . ?opera base:creatoDa ?creatore . }}
FINE REGOLE

### Conversazione ###
Utente: che ore sono?
Assistente:
"""

def create_explanation_prompt(results_str: str) -> str:
    """Prompt per generare una spiegazione museale dei risultati SPARQL."""
    return f"""\
Ho ottenuto questi risultati SPARQL:
{results_str}
Ora fornisci una breve spiegazione museale (massimo ~10 righe), senza inventare oltre i risultati.
"""

async def call_hf_model(prompt: str, temperature: float = 0.5, max_tokens: int = 150) -> str:
    """Chiama il modello Hugging Face tramite l'API REST e gestisce la risposta."""
    logger.debug("Chiamo HF con il seguente prompt:")
    content_preview = (prompt[:300] + '...') if len(prompt) > 300 else prompt
    logger.debug(f"PROMPT => {content_preview}")

    headers = {
        "Authorization": f"Bearer {API_KEY}"
    }
    payload = {
        "inputs": prompt,
        "parameters": {
            "temperature": temperature,
            "max_new_tokens": max_tokens,
            "top_p": 0.9
        }
    }

    try:
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{HF_MODEL}",
            headers=headers,
            json=payload
        )
        if response.status_code != 200:
            logger.error(f"Errore nella chiamata all'API Hugging Face: {response.status_code} - {response.text}")
            raise HTTPException(status_code=500, detail=f"Errore nell'API Hugging Face: {response.text}")
        data = response.json()
        logger.debug(f"Risposta completa dal modello: {data}")
        if isinstance(data, list) and len(data) > 0 and "generated_text" in data[0]:
            raw = data[0]["generated_text"]
        elif "generated_text" in data:
            raw = data["generated_text"]
        else:
            raise ValueError("Nessun campo 'generated_text' nella risposta.")

        # Forza la risposta su una singola linea se multilinea
        single_line = " ".join(raw.splitlines())
        logger.debug(f"Risposta HF single-line: {single_line}")
        return single_line.strip()
    except Exception as e:
        logger.error(f"Errore nella chiamata all'API Hugging Face tramite requests: {e}")
        raise HTTPException(status_code=500, detail=str(e))

# Variabile globale per le etichette delle entità
entity_labels: List[str] = []

def load_entity_labels(documents_file: str):
    """Carica le etichette delle entità dal file documents.json."""
    global entity_labels
    try:
        with open(documents_file, "r", encoding="utf-8") as f:
            document = json.load(f)
        entity_labels = [entity['label'].lower() for entity in document['entities']]
        logger.info(f"Elenco delle etichette delle entità caricato: {entity_labels}")
    except Exception as e:
        logger.error(f"Errore nel caricamento delle etichette delle entità: {e}")
        entity_labels = []

def is_ontology_related(query: str) -> bool:
    """Determina se la domanda è pertinente all'ontologia."""
    query_lower = query.lower()
    keywords = ["opera", "museo", "stanza", "tour", "visitatore", "biglietto", "guida", "evento", "agente"]
    if any(keyword in query_lower for keyword in keywords):
        return True
    if any(entity in query_lower for entity in entity_labels):
        return True
    return False

# Prepara i file necessari per RAG
prepare_retrieval()

# Carica il 'sunto' di classi, proprietà ed entità
knowledge_text = extract_classes_and_properties(RDF_FILE)

# Carica le etichette delle entità
load_entity_labels(DOCUMENTS_FILE)

app = FastAPI()

class QueryRequest(BaseModel):
    message: str
    max_tokens: int = 150  # Ridotto per risposte concise
    temperature: float = 0.5

@app.post("/generate-response/")
async def generate_response(req: QueryRequest):
    user_input = req.message
    logger.info(f"Utente dice: {user_input}")

    if not is_ontology_related(user_input):
        return {
            "type": "NATURAL",
            "response": "Ciao! Sono un assistente museale e non ho informazioni sulle ore attuali. Ti consiglio di consultare un orologio o un dispositivo mobile per conoscere l'ora esatta."
        }

    try:
        # Recupera documenti rilevanti usando RAG
        retrieved_docs = retrieve_relevant_documents(user_input, top_k=3)
    except Exception as e:
        logger.error(f"Errore nel recupero dei documenti rilevanti: {e}")
        return {"type": "ERROR", "response": f"Errore nel recupero dei documenti: {e}"}

    sys_msg = create_system_message(knowledge_text, retrieved_docs)
    prompt = f"{sys_msg}\nUtente: {user_input}\nAssistente:"

    # Primo tentativo
    try:
        r1 = await call_hf_model(prompt, req.temperature, req.max_tokens)
        logger.info(f"PRIMA RISPOSTA:\n{r1}")
    except Exception as e:
        logger.error(f"Errore nella chiamata al modello Hugging Face: {e}")
        return {"type": "ERROR", "response": f"Errore nella generazione della risposta: {e}"}

    # Se non parte con "PREFIX base:"
    if not r1.startswith("PREFIX base:"):
        sc = f"Non hai risposto con query SPARQL su una sola riga. Riprova. Domanda: {user_input}"
        fallback_prompt = f"{sys_msg}\nAssistente: {r1}\nUtente: {sc}\nAssistente:"
        try:
            r2 = await call_hf_model(fallback_prompt, req.temperature, req.max_tokens)
            logger.info(f"SECONDA RISPOSTA:\n{r2}")
            if r2.startswith("PREFIX base:"):
                sparql_query = r2
            else:
                return {"type": "NATURAL", "response": r2}
        except Exception as e:
            logger.error(f"Errore nella seconda chiamata al modello Hugging Face: {e}")
            return {"type": "ERROR", "response": f"Errore nella generazione della seconda risposta: {e}"}
    else:
        sparql_query = r1

    # Esegui la query con rdflib
    g = rdflib.Graph()
    try:
        g.parse(RDF_FILE, format="xml")
        logger.info(f"Parsing RDF di {RDF_FILE} riuscito per l'esecuzione della query.")
    except Exception as e:
        logger.error(f"Parsing RDF error: {e}")
        return {"type": "ERROR", "response": f"Parsing RDF error: {e}"}

    try:
        results = g.query(sparql_query)
        logger.info(f"Query SPARQL eseguita con successo. Risultati: {len(results)}")
    except Exception as e:
        fallback = f"La query SPARQL ha fallito. Riprova. Domanda: {user_input}"
        fallback_prompt = f"{sys_msg}\nAssistente: {sparql_query}\nUtente: {fallback}\nAssistente:"
        try:
            r3 = await call_hf_model(fallback_prompt, req.temperature, req.max_tokens)
            logger.info(f"TERZA RISPOSTA (fallback):\n{r3}")
            if r3.startswith("PREFIX base:"):
                sparql_query = r3
                try:
                    results = g.query(sparql_query)
                    logger.info(f"Seconda query SPARQL eseguita con successo. Risultati: {len(results)}")
                except Exception as e2:
                    logger.error(f"Seconda Query fallita: {e2}")
                    return {"type": "ERROR", "response": f"Query fallita di nuovo: {e2}"}
            else:
                return {"type": "NATURAL", "response": r3}
        except Exception as e:
            logger.error(f"Errore nella chiamata al modello Hugging Face durante il fallback: {e}")
            return {"type": "ERROR", "response": f"Errore durante il fallback della risposta: {e}"}

    if len(results) == 0:
        return {"type": "NATURAL", "sparql_query": sparql_query, "response": "Nessun risultato."}

    # Confeziona risultati
    row_list = []
    for row in results:
        row_dict = row.asdict()
        row_str = ", ".join([f"{k}: {v}" for k, v in row_dict.items()])
        row_list.append(row_str)
    results_str = "\n".join(row_list)

    # Spiegazione
    exp_prompt = create_explanation_prompt(results_str)
    try:
        explanation = await call_hf_model(exp_prompt, req.temperature, req.max_tokens)
    except Exception as e:
        logger.error(f"Errore nella generazione della spiegazione: {e}")
        return {"type": "ERROR", "response": f"Errore nella generazione della spiegazione: {e}"}

    return {
        "type": "NATURAL",
        "sparql_query": sparql_query,
        "sparql_results": row_list,
        "explanation": explanation
    }

@app.get("/")
def home():
    return {"message": "Assistente Museale con supporto SPARQL."}