import pdfplumber
import gradio as gr
import torch

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM
)

MODEL_NAME = "microsoft/Phi-3.5-mini-instruct"

print("Loading model...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    trust_remote_code=True,
    low_cpu_mem_usage=True
)

print("Model loaded successfully.")


def extract_text(pdf_file):
    text = ""

    try:
        with pdfplumber.open(pdf_file.name) as pdf:

            for page in pdf.pages:
                page_text = page.extract_text()

                if page_text:
                    text += page_text + "\n"

    except Exception as e:
        return f"PDF Extraction Error: {str(e)}"

    print(f"Extracted {len(text)} characters")

    # Keep small for free-tier inference
    return text[:1000]


def build_prompt(policy_text):

    return f"""
You are a senior insurance consultant.

Analyze the insurance policy and create a customer-friendly report.

Return markdown.

# Executive Summary

Summarize the policy in plain English.

# Customer Risk Score

Rate 1-10 and explain why.

# Policy Complexity Score

Rate 1-10 and explain why.

# Claim Difficulty Score

Rate 1-10 and explain why.

# What Is Covered

Provide bullet points.

# Major Exclusions

Provide bullet points.

# Waiting Periods

Provide bullet points.

# Coverage Gaps

Identify situations where customers may wrongly assume they are covered.

# Claim Checklist

Provide step-by-step instructions.

# Questions To Ask The Insurer

Provide 5 questions.

# Explain Like I'm 15

Explain the policy simply.

POLICY DOCUMENT:

{policy_text}
"""


def generate_response(prompt):

    messages = [
        {
            "role": "system",
            "content": "You are an expert insurance policy analyst."
        },
        {
            "role": "user",
            "content": prompt
        }
    ]

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=4096
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"

    model.to(device)

    
    inputs = {k: v.to(device) for k, v in inputs.items()}

    outputs = model.generate(
        **inputs,
        max_new_tokens=800,
        temperature=0.2,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id
    )

    generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]

    response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True
    )

    return response


def analyze_policy(pdf_file):

    try:

        if pdf_file is None:
            return "Please upload a policy PDF."

        policy_text = extract_text(pdf_file)

        if len(policy_text.strip()) == 0:
            return "No text could be extracted from this PDF."

        prompt = build_prompt(policy_text)

        response = generate_response(prompt)

        return response

    except Exception as e:

        print("ERROR:", e)

        return f"""
# Error

{str(e)}
"""


CUSTOM_CSS = """
footer {
    display:none;
}

.gradio-container {
    max-width: 1200px !important;
}
"""


with gr.Blocks(
    title="Insurance Policy Decoder",
    theme=gr.themes.Soft(),
    css=CUSTOM_CSS
) as demo:

    gr.Markdown(
        """
# 🛡️ Insurance Policy Decoder

Understand your insurance policy in less than a minute.

Upload a policy PDF and receive:

✅ Executive Summary

✅ Coverage Details

✅ Exclusions

✅ Waiting Periods

✅ Coverage Gaps

✅ Risk Scores

✅ Claim Checklist

✅ Questions To Ask Your Insurer
"""
    )

    pdf_input = gr.File(
        label="Upload Insurance Policy PDF",
        file_types=[".pdf"]
    )

    analyze_btn = gr.Button(
        "Decode Policy",
        variant="primary"
    )

    output = gr.Markdown(
        value="Upload a policy document and click **Decode Policy**."
    )

    analyze_btn.click(
        fn=analyze_policy,
        inputs=pdf_input,
        outputs=output,
        show_progress="full"
    )

demo.launch()