import gradio as gr
import requests
from datasets import load_dataset
from transformers import pipeline

# Load the dataset
dataset = load_dataset("viber1/indian-law-dataset")['train']

# Load a pre-trained language model for question-answering
qa_model = pipeline("question-answering", model="deepset/roberta-base-squad2")

def get_answer_from_api(query):
    # Use CourtListener API to get legal information
    base_url = "https://www.courtlistener.com/api/rest/v4/search/"
    headers = {
        "Authorization": "Token 9c70738ed9eb3cce4f3782a91c7c8a218c180b89"  # Replace with your actual API token
    }
    params = {
        "q": query,
        "page_size": 1  # Limit the number of results returned
    }

    try:
        response = requests.get(base_url, headers=headers, params=params)
        response.raise_for_status()  # Raise an error for bad responses
        results = response.json()

        # Check if there are any results
        if results.get('count', 0) > 0:
            return results['results'][0]['case_name']  # Adjust based on actual response structure
        else:
            return None  # No results found
    except requests.RequestException as e:
        print(f"API request failed: {e}")  # Print the error message for debugging
        return None  # Return None if there was an error

def get_answer_from_dataset(query):
    # Look for an answer in the dataset
    for entry in dataset:
        if query.lower() in entry['Instruction'].lower():
            return entry['Response']
    return None  # No answer found in the dataset

def get_answer_from_model(query):
    # Use the pre-trained model to generate an answer
    context = " ".join([entry['Response'] for entry in dataset])  # Combine all responses from dataset
    result = qa_model(question=query, context=context)
    return result['answer'] if result['score'] > 0.2 else None  # eturn answer if confidence score is high

def respond(query):
    # First, try to get the answer from the API
    answer = get_answer_from_dataset(query)
    if answer:
        return answer  # Return if found in API
    
    # If not found, look in the dataset
    answer = get_answer_from_model(query)
    if answer:
        return answer  # Return if found in dataset
    
    # If still no answer, use the model
    return get_answer_from_api(query)

# Gradio interface
demo = gr.Interface(
    fn=respond,
    inputs="text",
    outputs="text",
    title="AI Legal Assistant",
    description="Ask your legal queries regarding Indian laws"
)

if _name_ == "_main_":
    demo.launch()