import os
import requests
import json 
from pypdf import PdfReader
from langchain_community.llms import Ollama
from langchain_community.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.docstore.document import Document

# Define the OllamaEmbeddings class
class OllamaEmbeddings(Embeddings):
    def __init__(self, model='llama3.2', url='http://localhost:11434/api/embeddings'):
        self.model = model
        self.url = url

    def embed(self, text):
        headers = {
            'Content-Type': 'application/json'
        }
        payload = {
            'model': self.model,
            'prompt': text
        }
        
        response = requests.post(self.url, headers=headers, data=json.dumps(payload))
        
        if response.status_code == 200:
            response_json = response.json()
            # Print the response structure for debugging
            print("Response JSON:", 1) #response_json)
            # Check the actual structure of the response and extract embeddings accordingly
            if 'embedding' in response_json:
                return response_json['embedding']
            else:
                print("Key 'embedding' not found in response")
                return None
        else:
            print(f"Error {response.status_code}: {response.text}")
            raise Exception(f"Error {response.status_code}: {response.text}")

    def embed_documents(self, texts):
        return [self.embed(text) for text in texts]

    def embed_query(self, text):
        return self.embed(text)


# Loading the LLM
llm = Ollama(
    model="llama3.2",
    temperature=0.1
    # server="http://localhost:11434"  # Specify the server address
)

# Loading the document using PyPDF2
def load_pdf(file_path, num_pages=None):
    reader = PdfReader(file_path)
    text = ""
    pages = reader.pages[:num_pages] if num_pages else reader.pages
    for page in pages:
        text += page.extract_text() + "\n"
    return text

pdf_path = "tcsbancs-overarching-brochure-2023.pdf"
document_text = load_pdf(pdf_path, num_pages=5)

# Split the document into chunks
text_splitter = CharacterTextSplitter(separator="\n",
                                      chunk_size=256,
                                      chunk_overlap=32)

text_chunks = text_splitter.split_text(document_text)

# Convert text chunks to document objects
documents = [Document(page_content=chunk) for chunk in text_chunks]




# Loading the vector embedding model
embeddings = OllamaEmbeddings()



# Create knowledge base
knowledge_base = FAISS.from_documents(documents, embeddings)



# Retrieval QA chain
qa_chain = RetrievalQA.from_chain_type(
    llm,
    retriever=knowledge_base.as_retriever()
)



question = "What is this document about?"
response = qa_chain.invoke({"query": question})
print(response["result"])

### Sample Prompt for Summarization and Extraction 


question_prompt = f"""
    Based on the document provided, perform the following tasks.

    Task:
    1. Summarize the document by providing a concise summary of the main points, focusing on the features, functionalities, \
    and benefits of the TCS BaNCS solution.
    2. Extract all the important information that describe how the solution has helped various banks like SBI and influenced industries like Capital \
    Markets in foreign countries. Extract as much information as possible.

    Provide the summary first, followed by the extracted information.

    Summary:
    """
response = qa_chain.invoke({"query": question_prompt})
print(response["result"])

### Converting Summary and Extracted Metrics into HTML format (Sample Task Completion Scenario)

question_prompt = f"""
    Based on the document provided, perform the following tasks.

    Task:
    1. Summarize the document by providing a concise summary of the main points, focusing on the features,\
    functionalities, and benefits of the TCS BaNCS solution.
    2. Extract all the important information that describe how the solution has helped various banks like SBI \
    and influenced industries like Capital Markets in foreign countries. Extract as much information as possible. \
    Provide the summary first, followed by the extracted information.
    3. Format everything as HTML content which can be used to build a website. Provide the title at the top \
    of the webpage as "TCS BaNCS Overview". Place the summary in <div> tags.
    4. Create a table titled Key Metrics and place the information extracted as rows of data. Use bold tags \
    to highlight crucial elements like names and numbers. 
    
    HTML Code:
    """
response = qa_chain.invoke({"query": question_prompt})
print(response["result"])

### Load Python libraries to view HTML

from IPython.display import display, HTML
display(HTML(response['result']))