import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import json 
import os
import re


hf_token = os.getenv("HK_TOKEN")  # Fetch Hugging Face API token

# Define model name and authentication token
model_name = "meta-llama/Llama-2-7b-chat-hf"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype=torch.float16, device_map="auto", use_auth_token=hf_token)


def compare_claims(claim_json1, claim_json2):
    """
    Compare two health insurance claims using OpenAI GPT model and return structured JSON output.

    :param claim_json1: JSON string for first claim
    :param claim_json2: JSON string for second claim
    :return: JSON output with similarities, differences, and summary
    """
    # Parse input JSONs
    claim1 = json.loads(claim_json1)
    claim2 = json.loads(claim_json2)

  # Structured prompt for better JSON output
    prompt = f"""
    Compare the following insurance claims and identify similarities and differences.
    
    Claim 1:
    {json.dumps(claim1, indent=2)}
    
    Claim 2:
    {json.dumps(claim2, indent=2)}
    
    Provide a structured comparison highlighting commonalities and differences.
    """

    device = "cuda" if torch.cuda.is_available() else "cpu"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
     

    #####################JSON handling
    output = model.generate(input_ids, max_length=512)
    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)

    # Extract only the JSON part using regex
    json_match = re.search(r"\{.*\}", decoded_output, re.DOTALL)
    if json_match:
        json_output = json_match.group(0)
        try:
            parsed_json = json.loads(json_output)
            print("Valid JSON Output:", parsed_json)
        except json.JSONDecodeError:
            print("Error: Model did not return valid JSON. Trying to fix it.")
            # Attempt a fix using simple corrections (e.g., missing brackets)
            json_output = json_output.rstrip(",")  # Remove trailing commas
            json_output = json_output.replace("\n", "")  # Remove new lines
            try:
                parsed_json = json.loads(json_output)
                print("Fixed JSON Output:", parsed_json)
            except json.JSONDecodeError:
                print("Final Error: Could not fix JSON formatting.")
    else:
        print("No JSON detected in model output.")

    #########################################
    
    # Generate Response
    output = model.generate(input_ids, max_length=200, do_sample=True, temperature=0.7)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    ##########################
    # Extract JSON part
    match = re.search(r"\{.*\}", response, re.DOTALL)
    if match:
        json_str = match.group(0)
        try:
            parsed_json = json.loads(json_str)
            print("Valid JSON Output:", parsed_json)
            return json.dumps(parsed_json, indent=4)
        except json.JSONDecodeError:
            print("Error: Invalid JSON returned by LLaMA")
    else:
        print("No JSON found in response.")
        
    ############################
    return




def main():
    """
    Launch the Gradio interface for claim comparison.
    """
    # Define the Gradio interface
    interface = gr.Interface(
        fn=compare_claims,
        inputs=[
            gr.Textbox(label="claim_json1", placeholder="Enter first claim description..."),
            gr.Textbox(label="claim_json2", placeholder="Enter second claim description...")
        ],
        outputs="text",
        title="Claims Comparison",
        description="Enter two claims to compare their differences."
    )

    # Launch the Gradio app
    interface.launch()

if __name__ == "__main__":
    main()