app/services/risk_analyzer.py

import google.generativeai as genai
from app.config import settings
from app.services.risk_scorer import RISK_DEFINITIONS
import json
import re
from typing import Dict


# Configure Gemini API
genai.configure(api_key=settings.GEMINI_API_KEY)


def analyze_clause_with_gemini(clause_text: str) -> Dict:
    """
    Analyze a contract (clause or full text) using Google Gemini AI for risk identification.
    Args:
        clause_text: The text of the clause or contract to analyze
    Returns:
        Dictionary containing identified risk IDs and suggestions
    """

    prompt = f"""
You are an expert Indian legal consultant specializing in contract analysis and risk assessment. 
Analyze the following contract text and identify any legal risks based on the predefined categories.

CONTRACT TEXT TO ANALYZE:
{clause_text}

RISK CATEGORIES TO CHECK FOR:
1. UNLIMITED_LIABILITY: Clause imposes unlimited liability on the client
2. ONE_SIDED_TERMINATION: Termination rights are unfairly one-sided
3. UNCLEAR_JURISDICTION: Governing law or jurisdiction for disputes is ambiguous
4. DPDP_NON_COMPLIANCE: Data protection clause may not comply with the DPDP Act 2023

INSTRUCTIONS:
1. Carefully read the contract text
2. Identify which of the above risk categories apply
3. For each identified risk, provide a **short explanation (max 1–2 lines, no long paragraphs)**
4. Suggest a **brief compliant alternative (1–2 lines)**
5. If no risks are found, respond with "No risks identified"

RESPONSE FORMAT (strict JSON only):
{{
    "risks": [
        {{
            "risk_id": "RISK_CATEGORY_ID",
            "explanation": "1–2 line explanation of why this risk applies"
        }}
    ],
    "suggestion": "1–2 line compliant alternative or modification suggestion"
}}

If no risks are identified, return:
{{
    "risks": [],
    "suggestion": "No risks identified - contract appears compliant"
}}
"""

    try:
        # Initialize Gemini model
        model = genai.GenerativeModel("gemini-2.5-flash-lite")

        # Generate response
        response = model.generate_content(prompt)
        response_text = response.text.strip()

        # Try to parse JSON from the response
        try:
            json_match = re.search(r"\{.*\}", response_text, re.DOTALL)
            if json_match:
                json_str = json_match.group()
                result = json.loads(json_str)
            else:
                result = json.loads(response_text)
        except json.JSONDecodeError:
            result = {
                "risks": [],
                "suggestion": "Unable to parse AI response - manual review recommended"
            }

        # Ensure required keys exist
        if "risks" not in result:
            result["risks"] = []
        if "suggestion" not in result:
            result["suggestion"] = "No suggestion provided"

        # Validate risk IDs against known definitions
        valid_risks = []
        for risk in result["risks"]:
            if isinstance(risk, dict) and "risk_id" in risk:
                risk_id = risk["risk_id"]
                if risk_id in RISK_DEFINITIONS:
                    valid_risks.append(risk)

        result["risks"] = valid_risks
        return result

    except Exception as e:
        print(f"Error in Gemini analysis: {e}")
        return {
            "risks": [],
            "suggestion": f"Analysis failed: {str(e)}"
        }