import PyPDF2
import re
from sentence_transformers import SentenceTransformer
import faiss
from langchain.agents import initialize_agent, AgentType,Tool
from langchain.schema import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
import gradio as gr
import os
import pytesseract
from PIL import Image
pytesseract.pytesseract.tesseract_cmd = r"tesseract.exe"
def retrieve_info():
    return """IPC in gist:
Title and Scope: The Indian Penal Code applies to the whole of India, except for the State of Jammu and Kashmir. It dictates punishments for acts or omissions that violate its provisions.●
General Explanations: The document provides definitions for various terms used throughout the code, such as "public," "servant of Government," "Government," "India," "Judge," "Court of Justice," and "moveable property."●
Punishments: The Indian Penal Code outlines six categories of punishment: death, imprisonment for life, imprisonment (rigorous or simple), forfeiture of property, and fine.●
Offenses Against the State: These include waging war against the Government of India, sedition, and offenses relating to the Army, Navy, and Air Force.●
Offenses Against Public Tranquility: This category encompasses unlawful assembly, rioting, and affray (fighting in a public place).●
Offenses By or Relating to Public Servants: This section covers offenses such as public servants disobeying the law, engaging in unlawful trade, and framing incorrect documents.●
Offenses Relating to Documents and Property Marks: This part deals with forgery, using forged documents, counterfeiting, and offenses related to property marks.●
Criminal Breach of Contracts of Service: The code addresses breaches of contract related to attending to helpless individuals.●
Offenses Relating to Marriage: This section covers offenses such as causing a woman to cohabit based on deceitful marriage, marrying again while a spouse is alive, and concealing a previous marriage.●
Defamation: The code defines defamation and includes exceptions for publications made in good faith, such as court reports, opinions on decided cases, and cautions conveyed for the good of others.●
Criminal Intimidation, Insult, and Annoyance: This section addresses criminal intimidation, intentional insult to provoke breach of peace, statements conducing to public mischief, and statements promoting enmity between classes.
"""



ipc_tool = Tool(
    name="IPC Information Retrieval",
    func=retrieve_info,
    description="Provides IPC Gist.No parameter required."
)


llm = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",
    temperature=0.25,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    prompt_template="""
    You are a highly specialized legal assistant with deep knowledge of the Indian Penal Code (IPC). 
    Your primary task is to retrieve and summarize legal information accurately from the IPC.pdf document provided to you. 
    Your responses should be highly specific, fact-based, and free from any speculation or hallucinations. 
    Always cite the exact section from the IPC when providing an answer. 
    If the information is not available in the document, clearly state that and do not make any assumptions.

    Example task: "What is the punishment for theft according to the IPC?"
    Example response: "According to Section 379 of the IPC, the punishment for theft is imprisonment of either description for a term which may extend to three years, or with fine, or with both."

    Task: {{query}}

    Response:
    """,
)

agent_tools = [ipc_tool]

agent = initialize_agent(
    tools=agent_tools,
    llm=llm,
    agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    return_intermediate_steps=True,
    handle_parsing_errors=True,
)
def encode_image_to_base64(image_path):
    return pytesseract.image_to_string(Image.open(image_path))
def chatbot_response(query):
    if query.get('files'):
        # Encode image to base64
        image_data=""
        for x in range(len(query["files"])):
            image_data += f"{x}. "+encode_image_to_base64(query["files"][x]) +"\n"
        
        # Create a multimodal message with both text and image data
        message = HumanMessage(
            content=[
                {"type": "text", "text": query['text'] +" System :Image(s) was added to this prompt by this user. Text Extracted from this image (Some words may be misspelled ,Use your understanding ):"+image_data},  # Add text input
               
            ]
        )
    else:
        # If no image, only pass the text
        message = HumanMessage(content=[{"type": "text", "text": query}])

    # Invoke the model with the multimodal message
    result = agent.invoke([message])
    response = result['output']
    intermediate_steps = result.get('intermediate_steps', [])
    
    thought_process = ""
    for action, observation in intermediate_steps:
        thought_process += f"**Thought:** {action.log}\n"
        thought_process += f"**Action:** {action.tool}\n"
        thought_process += f"**Observation:** {observation}\n\n"

    return response, thought_process.strip()
# Step 5: Gradio Interface
from gradio import ChatMessage
def chatbot_interface(messages,prompt):
    response, thought_process = chatbot_response(prompt)
    #messages.append(ChatMessage(role="user", content=prompt))
    
    for x in prompt["files"]:
            messages.append(ChatMessage(role="user", content={"path": x, "mime_type": "image/png"}))
    if prompt["text"] is not None:
            messages.append(ChatMessage(role="user", content=prompt['text']))
    if thought_process:
        messages.append(ChatMessage(role="assistant", content=thought_process,metadata={"title": "🧠 Thought Process"}))
    messages.append(ChatMessage(role="assistant", content=response))
   
    return messages,  gr.MultimodalTextbox(value=None, interactive=True)


def vote(data: gr.LikeData):
    if data.liked:
        print("You upvoted this response: " + data.value)
    else:
        print("You downvoted this response: " + data.value)

with gr.Blocks(theme=gr.themes.Soft()) as iface:
   
            gr.Markdown(
                """
                <div style="font-size: 24px; font-weight: bold; color: #333;">
                    DoJ Chatbot
                </div>
                <div style="font-size: 16px; color: #555;">
                    Ask questions related to the Department of Justice.
                </div>
                """
            )
            chatbot = gr.Chatbot(type="messages",avatar_images=("user.jpeg", "logo.jpeg"), bubble_full_width=True)  # Chatbot component to display conversation history
            query_input = gr.MultimodalTextbox(interactive=True,
                                      placeholder="Enter message or upload file...", show_label=False)
            submit_button = gr.Button("Send")

            submit_button.click(chatbot_interface, [chatbot, query_input], [chatbot, query_input])
            query_input.submit(chatbot_interface, [chatbot, query_input], [chatbot,query_input])

            chatbot.like(vote, None, None)  # Adding like/dislike functionality to the chatbot

        
iface.launch(
    show_error=True,
    prevent_thread_lock=True
)