from transformers import pipeline
import gradio as gr

# Load pre-trained pipelines
try:
    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
    ner = pipeline("ner", model="Davlan/bert-base-multilingual-cased-ner-hrl", aggregation_strategy="simple")
except Exception as e:
    summarizer = None
    ner = None
    print("Error loading models:", e)

# Nigerian law reference (basic keyword-to-punishment mapping)
crime_punishment_map = {
    "theft": {"law": "Criminal Code Act, Section 390", "punishment": "Up to 3 years imprisonment"},
    "armed robbery": {"law": "Robbery and Firearms Act, Section 1", "punishment": "Death penalty or life imprisonment"},
    "internet fraud": {"law": "Cybercrime Act, 2015", "punishment": "Minimum of 7 years imprisonment"},
    "rape": {"law": "Criminal Law of Lagos State, Section 260", "punishment": "Life imprisonment"},
    "murder": {"law": "Criminal Code Act, Section 319", "punishment": "Death penalty"},
    "assault": {"law": "Criminal Code Act, Section 351", "punishment": "1 year imprisonment"}
}

def classify_crime(text):
    text = text.lower()
    for crime in crime_punishment_map:
        if crime in text:
            return crime, crime_punishment_map[crime]
    return "unknown", {
        "law": "N/A",
        "punishment": "No specific punishment found. Manual review required."
    }

# Main analysis function with full error handling
def analyze_text(text):
    try:
        if not text.strip():
            return "No text provided.", [], {"Crime Type": "N/A", "Applicable Law": "N/A", "Recommended Punishment": "N/A"}

        summary = summarizer(text, max_length=80, min_length=30, do_sample=False)[0].get("summary_text", "Summary failed.")
        entities = ner(text)
        crime_type, law_info = classify_crime(text)

        return summary, entities, {
            "Crime Type": crime_type.title() if crime_type != "unknown" else "Unknown",
            "Applicable Law": law_info["law"],
            "Recommended Punishment": law_info["punishment"]
        }
    except Exception as e:
        return f"⚠️ An error occurred: {str(e)}", [], {
            "Crime Type": "Error",
            "Applicable Law": "Error",
            "Recommended Punishment": "Error"
        }

# Launch app
gr.Interface(
    fn=analyze_text,
    inputs=gr.Textbox(lines=12, label="Paste Criminal Case Text"),
    outputs=[
        gr.Textbox(label="Summary"),
        gr.JSON(label="Extracted Entities"),
        gr.JSON(label="Legal Analysis / Recommended Punishment")
    ],
    title="JusticeAI - Legal Case Analyzer",
    description="Paste any criminal case report. This AI will summarize it, extract important entities, and recommend the legal punishment based on Nigerian law."
).launch()