from flask import Flask, request
from twilio.twiml.messaging_response import MessagingResponse
from twilio.rest import Client
import os
import requests
from PIL import Image
import shutil

from langchain.vectorstores.chroma import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_community.llms.ollama import Ollama
from get_embedding_function import get_embedding_function
from langchain.document_loaders.pdf import PyPDFDirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.schema.document import Document

app = Flask(__name__)
UPLOAD_FOLDER = '/code/uploads'
CHROMA_PATH = UPLOAD_FOLDER  # Use the same folder for Chroma
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

class ConversationBufferMemory:
    def __init__(self, max_size=6):
        self.memory = []
        self.max_size = max_size

    def add_to_memory(self, interaction):
        self.memory.append(interaction)
        if len(self.memory) > self.max_size:
            self.memory.pop(0)

    def get_memory(self):
        return self.memory

conversation_memory = ConversationBufferMemory(max_size=2)

account_sid = os.environ.get('TWILIO_ACCOUNT_SID')
auth_token = os.environ.get('TWILIO_AUTH_TOKEN')
client = Client(account_sid, auth_token)
from_whatsapp_number = 'whatsapp:+14155238886'

PROMPT_TEMPLATE = """
Answer the question based only on the following context:
{context}
---
Answer the question based on the above context: {question}
"""

from bs4 import BeautifulSoup
import requests
from requests.auth import HTTPBasicAuth
from PIL import Image
from io import BytesIO
import pandas as pd
from urllib.parse import urlparse
import os
from pypdf import PdfReader
from ai71 import AI71
import uuid

from inference_sdk import InferenceHTTPClient
import base64

AI71_API_KEY = os.environ.get('AI71_API_KEY')

def generate_response(query, chat_history):
    response = ''
    for chunk in AI71(AI71_API_KEY).chat.completions.create(
            model="tiiuae/falcon-180b-chat",
            messages=[
                {"role": "system", "content": "You are the best agricultural assistant. Remember to give a response in not more than 2 sentences. Greet the user if the user greets you."},
                {"role": "user", "content": f'''Answer the query based on history {chat_history}: {query}'''},
            ],
            stream=True,
    ):
        if chunk.choices[0].delta.content:
            response += chunk.choices[0].delta.content
    return response.replace("###", '').replace('\nUser:', '')

def predict_pest(filepath):
    CLIENT = InferenceHTTPClient(
        api_url="https://detect.roboflow.com",
        api_key="oF1aC4b1FBCDtK8CoKx7"
    )
    result = CLIENT.infer(filepath, model_id="pest-detection-ueoco/1")
    return result['predictions'][0]
    

def predict_disease(filepath):
    CLIENT = InferenceHTTPClient(
        api_url="https://classify.roboflow.com",
        api_key="oF1aC4b1FBCDtK8CoKx7"
    )
    result = CLIENT.infer(filepath, model_id="plant-disease-detection-iefbi/1")
    return result['predicted_classes'][0]

def convert_img(url, account_sid, auth_token):
    try:
        response = requests.get(url, auth=HTTPBasicAuth(account_sid, auth_token))
        response.raise_for_status()

        parsed_url = urlparse(url)
        media_id = parsed_url.path.split('/')[-1]
        filename = f"downloaded_media_{media_id}"

        media_filepath = os.path.join(UPLOAD_FOLDER, filename)
        with open(media_filepath, 'wb') as file:
            file.write(response.content)
        
        print(f"Media downloaded successfully and saved as {media_filepath}")

        with open(media_filepath, 'rb') as img_file:
            image = Image.open(img_file)

            converted_filename = f"image.jpg"
            converted_filepath = os.path.join(UPLOAD_FOLDER, converted_filename)
            image.convert('RGB').save(converted_filepath, 'JPEG')
            return converted_filepath

    except requests.exceptions.HTTPError as err:
        print(f"HTTP error occurred: {err}")
    except Exception as err:
        print(f"An error occurred: {err}")

def get_weather(city):
    city = city.strip().replace(' ', '+')
    r = requests.get(f'https://www.google.com/search?q=weather+in+{city}')
    soup = BeautifulSoup(r.text, 'html.parser')
    temperature = soup.find('div', attrs={'class': 'BNeawe iBp4i AP7Wnd'}).text
    return temperature

from zenrows import ZenRowsClient
Zenrow_api = os.environ.get('Zenrow_api')
zenrows_client = ZenRowsClient(Zenrow_api)

def get_rates():
    url = "https://www.kisandeals.com/mandiprices/ALL/TAMIL-NADU/ALL"
    response = zenrows_client.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.content, 'html.parser')
        rows = soup.select('table tbody tr')
        data = {}
        for row in rows:
            columns = row.find_all('td')
            if len(columns) >= 2:
                commodity = columns[0].get_text(strip=True)
                price = columns[1].get_text(strip=True)
                if '₹' in price:
                    data[commodity] = price
    return str(data) + " These are the prices for 1 kg"

def get_news():
    news = []
    url = "https://economictimes.indiatimes.com/news/economy/agriculture?from=mdr"
    response = zenrows_client.get(url)

    if response.status_code == 200:
        soup = BeautifulSoup(response.content, 'html.parser')
        headlines = soup.find_all("div", class_="eachStory")
        for story in headlines:
            headline = story.find('h3').text.strip()
            news.append(headline)
    return news

def download_and_save_as_txt(url, account_sid, auth_token):
    try:
        response = requests.get(url, auth=HTTPBasicAuth(account_sid, auth_token))
        response.raise_for_status()

        parsed_url = urlparse(url)
        media_id = parsed_url.path.split('/')[-1]
        filename = f"pdf_file.pdf"

        txt_filepath = os.path.join(UPLOAD_FOLDER, filename)
        with open(txt_filepath, 'wb') as file:
            file.write(response.content)
        
        print(f"Media downloaded successfully and saved as {txt_filepath}")
        return txt_filepath

    except requests.exceptions.HTTPError as err:
        print(f"HTTP error occurred: {err}")
    except Exception as err:
        print(f"An error occurred: {err}")

def query_rag(query_text: str):
    embedding_function = get_embedding_function()
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embedding_function)
    results = db.similarity_search_with_score(query_text, k=5)
    context_text = "\n\n---\n\n".join([doc.page_content for doc, _score in results])
    prompt_template = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)
    prompt = prompt_template.format(context=context_text, question=query_text)
    model = Ollama(model="llama2")
    response_text = model.invoke(prompt)
    return response_text

def save_pdf_and_update_database(media_url):
    response = requests.get(media_url)
    pdf_filename = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4()}.pdf")
    with open(pdf_filename, 'wb') as f:
        f.write(response.content)
    
    document_loader = PyPDFDirectoryLoader(UPLOAD_FOLDER)
    documents = document_loader.load()
    
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=800,
        chunk_overlap=80,
        length_function=len,
        is_separator_regex=False,
    )
    chunks = text_splitter.split_documents(documents)
    
    add_to_chroma(chunks)

def add_to_chroma(chunks: list[Document]):
    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=get_embedding_function())
    chunks_with_ids = calculate_chunk_ids(chunks)
    existing_items = db.get(include=[])
    existing_ids = set(existing_items["ids"])

    new_chunks = [chunk for chunk in chunks_with_ids if chunk.metadata["id"] not in existing_ids]

    if new_chunks:
        new_chunk_ids = [chunk.metadata["id"] for chunk in new_chunks]
        db.add_documents(new_chunks, ids=new_chunk_ids)
        db.persist()

def calculate_chunk_ids(chunks):
    last_page_id = None
    current_chunk_index = 0

    for chunk in chunks:
        source = chunk.metadata.get("source")
        page = chunk.metadata.get("page")
        current_page_id = f"{source}:{page}"

        if current_page_id == last_page_id:
            current_chunk_index += 1
        else:
            current_chunk_index = 0

        last_page_id = current_page_id
        chunk_id = f"{current_page_id}:{current_chunk_index}"
        chunk.metadata["id"] = chunk_id

    return chunks

@app.route("/pdf", methods=["POST"])
def receive_pdf():
    media_url = request.values.get("MediaUrl", None)
    if media_url:
        save_pdf_and_update_database(media_url)
        return "PDF processed and saved successfully."
    return "No media URL found."

@app.route("/whatsapp", methods=["POST"])
def incoming_whatsapp():
    media_url = request.values.get("MediaUrl", None)
    from_number = request.values.get("From", "").strip()
    from_number = from_number[2:] if from_number.startswith("91") else from_number
    incoming_msg = request.values.get('Body', '').lower()
    response = MessagingResponse()
    message = response.message()

    if media_url:
        extension = os.path.splitext(media_url)[1]
        if extension.lower() == ".pdf":
            media_filepath = download_and_save_as_txt(media_url, account_sid, auth_token)
            save_pdf_and_update_database(media_url)
            message.body("The PDF was processed successfully.")
        else:
            message.body("Please send a PDF file.")
        return str(response)

    if 'get weather for' in incoming_msg:
        city = incoming_msg.replace("get weather for", "")
        temperature = get_weather(city)
        message.body(f'The temperature in {city} is {temperature}.')
        return str(response)

    if 'get rates' in incoming_msg:
        message.body(get_rates())
        return str(response)

    if 'get news' in incoming_msg:
        message.body(get_news())
        return str(response)

    if 'pest' in incoming_msg:
        text = predict_pest(media_filepath)
        message.body(text)
        return str(response)

    if 'disease' in incoming_msg:
        text = predict_disease(media_filepath)
        message.body(text)
        return str(response)

    if 'question:' in incoming_msg:
        conversation_memory.add_to_memory(f"User: {incoming_msg}")
        chat_history = "\n".join(conversation_memory.get_memory())
        response_text = generate_response(incoming_msg.replace("question:", ""), chat_history)
        conversation_memory.add_to_memory(f"Assistant: {response_text}")
        message.body(response_text)
        return str(response)

    if 'query:' in incoming_msg:
        query = incoming_msg.replace("query:", "").strip()
        response_text = query_rag(query)
        message.body(response_text)
        return str(response)

    message.body("I'm sorry, I don't understand that command.")
    return str(response)

if __name__ == "__main__":
    app.run(debug=True)