from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
import streamlit as st

class RAGSystem:
    def __init__(self):
        self.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-mpnet-base-v2"
        )
        self.vector_store = None
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50
        )
        
    def initialize_knowledge_base(self, knowledge_base):
        """Initialize vector store with construction knowledge"""
        documents = []
        
        for damage_type, cases in knowledge_base.items():
            for case in cases:
                doc_text = f"""
                Damage Type: {damage_type}
                Severity: {case['severity']}
                Description: {case['description']}
                Repair Methods: {', '.join(case['repair_method'])}
                Cost: {case['estimated_cost']}
                Timeframe: {case['timeframe']}
                Location: {case['location']}
                Required Expertise: {case['required_expertise']}
                Immediate Action: {case['immediate_action']}
                Prevention: {case['prevention']}
                """
                documents.append(doc_text)
        
        splits = self.text_splitter.create_documents(documents)
        self.vector_store = FAISS.from_documents(splits, self.embeddings)
        
    def query_knowledge_base(self, query, k=3):
        """Query the vector store for relevant information"""
        if not self.vector_store:
            raise ValueError("Vector store not initialized")
            
        results = self.vector_store.similarity_search(query, k=k)
        return results
    
    def get_enhanced_analysis(self, damage_type, confidence, custom_query=None):
        """Get enhanced analysis based on damage type and confidence"""
        if not custom_query:
            query = f"What are the key considerations and recommendations for {damage_type} with {confidence}% confidence?"
        else:
            query = custom_query
            
        relevant_docs = self.query_knowledge_base(query)
        
        enhanced_info = {
            "technical_details": [],
            "safety_considerations": [],
            "expert_recommendations": []
        }
        
        for doc in relevant_docs:
            content = doc.page_content
            
            if "severity" in content.lower() or "description" in content.lower():
                enhanced_info["technical_details"].append(content)
            elif "immediate action" in content.lower() or "prevention" in content.lower():
                enhanced_info["safety_considerations"].append(content)
            elif "repair" in content.lower() or "expertise" in content.lower():
                enhanced_info["expert_recommendations"].append(content)
                
        return enhanced_info