"""
Cata Risk Lab: Policy Auditor
A Streamlit tool for auditing AI Use Policies
"""

import streamlit as st
import re
from dataclasses import dataclass


@dataclass
class KeywordCheck:
    keyword: str
    found: bool
    weight: int
    category: str


def analyze_policy(text: str) -> dict:
    """Analyze the AI Use Policy text and return findings."""
    
    text_lower = text.lower()
    
    # Define keywords to check with their weights and categories
    keywords_config = [
        ("liability", 15, "Legal Protection"),
        ("human review", 15, "Oversight"),
        ("data training", 10, "Data Governance"),
        ("human-in-the-loop", 20, "Critical Safety"),
        ("accountability", 10, "Legal Protection"),
        ("transparency", 10, "Ethics"),
        ("bias", 8, "Fairness"),
        ("audit", 8, "Compliance"),
        ("consent", 7, "Privacy"),
        ("privacy", 7, "Privacy"),
        ("security", 5, "Security"),
        ("compliance", 5, "Compliance"),
    ]
    
    results = []
    total_possible = sum(k[1] for k in keywords_config)
    earned_points = 0
    
    for keyword, weight, category in keywords_config:
        # Check for keyword (allowing for variations)
        pattern = re.compile(re.escape(keyword), re.IGNORECASE)
        found = bool(pattern.search(text))
        
        if found:
            earned_points += weight
        
        results.append(KeywordCheck(
            keyword=keyword,
            found=found,
            weight=weight,
            category=category
        ))
    
    # Calculate base score (normalized to 100)
    base_score = (earned_points / total_possible) * 100
    
    # Special penalty: Deduct additional points if 'Human-in-the-Loop' is missing
    human_in_loop_check = next((r for r in results if r.keyword == "human-in-the-loop"), None)
    penalty = 0
    if human_in_loop_check and not human_in_loop_check.found:
        penalty = 15  # Additional penalty beyond the missed weight
    
    final_score = max(0, base_score - penalty)
    
    return {
        "results": results,
        "base_score": base_score,
        "penalty": penalty,
        "final_score": round(final_score, 1),
        "earned_points": earned_points,
        "total_possible": total_possible
    }


def get_score_color(score: float) -> str:
    """Return color based on score."""
    if score >= 80:
        return "#28a745"  # Green
    elif score >= 60:
        return "#ffc107"  # Yellow
    elif score >= 40:
        return "#fd7e14"  # Orange
    else:
        return "#dc3545"  # Red


def get_score_label(score: float) -> str:
    """Return label based on score."""
    if score >= 80:
        return "Excellent"
    elif score >= 60:
        return "Good"
    elif score >= 40:
        return "Needs Improvement"
    else:
        return "High Risk"


def render_badge():
    """Render the certification badge."""
    badge_html = """
    <div style="
        display: flex;
        justify-content: center;
        margin: 20px 0;
    ">
        <div style="
            background: linear-gradient(135deg, #1a5f2a 0%, #28a745 50%, #1a5f2a 100%);
            border: 3px solid #ffd700;
            border-radius: 15px;
            padding: 20px 40px;
            text-align: center;
            box-shadow: 0 4px 15px rgba(0,0,0,0.3);
        ">
            <div style="font-size: 40px; margin-bottom: 5px;">🛡️</div>
            <div style="
                color: #ffd700;
                font-size: 14px;
                font-weight: bold;
                letter-spacing: 2px;
                margin-bottom: 5px;
            ">✓ CERTIFIED</div>
            <div style="
                color: white;
                font-size: 18px;
                font-weight: bold;
            ">Cata Risk Lab</div>
            <div style="
                color: #90EE90;
                font-size: 12px;
                margin-top: 5px;
            ">AI Policy Approved</div>
        </div>
    </div>
    """
    st.markdown(badge_html, unsafe_allow_html=True)


def main():
    st.set_page_config(
        page_title="Cata Risk Lab: Policy Auditor",
        page_icon="🔍",
        layout="wide"
    )
    
    # Custom CSS
    st.markdown("""
        <style>
        .main-header {
            text-align: center;
            padding: 20px;
            background: linear-gradient(90deg, #1e3a5f, #2d5a87);
            border-radius: 10px;
            margin-bottom: 30px;
        }
        .main-header h1 {
            color: white;
            margin: 0;
        }
        .main-header p {
            color: #a0c4e8;
            margin: 5px 0 0 0;
        }
        .keyword-found {
            background-color: #d4edda;
            border-left: 4px solid #28a745;
            padding: 10px;
            margin: 5px 0;
            border-radius: 0 5px 5px 0;
        }
        .keyword-missing {
            background-color: #f8d7da;
            border-left: 4px solid #dc3545;
            padding: 10px;
            margin: 5px 0;
            border-radius: 0 5px 5px 0;
        }
        .score-card {
            text-align: center;
            padding: 30px;
            border-radius: 15px;
            margin: 20px 0;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Header
    st.markdown("""
        <div class="main-header">
            <h1>🔍 Cata Risk Lab: Policy Auditor</h1>
            <p>Analyze your AI Use Policy for safety and compliance</p>
        </div>
    """, unsafe_allow_html=True)
    
    # Main layout
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("📄 Paste Your AI Use Policy")
        
        policy_text = st.text_area(
            "Enter your company's AI Use Policy below:",
            height=400,
            placeholder="""Paste your AI Use Policy here...

Example content might include:
- Data handling procedures
- Human oversight requirements
- Liability clauses
- Training data policies
- Compliance frameworks"""
        )
        
        analyze_button = st.button("🔍 Analyze Policy", type="primary", use_container_width=True)
    
    with col2:
        st.subheader("📊 Analysis Results")
        
        if analyze_button and policy_text.strip():
            analysis = analyze_policy(policy_text)
            
            # Score display
            score = analysis["final_score"]
            score_color = get_score_color(score)
            score_label = get_score_label(score)
            
            st.markdown(f"""
                <div class="score-card" style="background: linear-gradient(135deg, {score_color}22, {score_color}44); border: 2px solid {score_color};">
                    <div style="font-size: 60px; font-weight: bold; color: {score_color};">{score}</div>
                    <div style="font-size: 20px; color: {score_color};">Safety Score / 100</div>
                    <div style="font-size: 16px; color: #666; margin-top: 10px;">{score_label}</div>
                </div>
            """, unsafe_allow_html=True)
            
            # Certification badge
            if score > 80:
                render_badge()
            else:
                st.warning("⚠️ Score must be above 80 to receive certification.")
            
            # Penalty notice
            if analysis["penalty"] > 0:
                st.error(f"🚨 **Penalty Applied:** -{analysis['penalty']} points for missing 'Human-in-the-Loop' provision")
            
            # Detailed results
            st.markdown("---")
            st.subheader("🔎 Detailed Findings")
            
            # Group by category
            categories = {}
            for result in analysis["results"]:
                if result.category not in categories:
                    categories[result.category] = []
                categories[result.category].append(result)
            
            for category, items in categories.items():
                with st.expander(f"📁 {category}", expanded=True):
                    for item in items:
                        if item.found:
                            st.markdown(f"""
                                <div class="keyword-found">
                                    ✅ <strong>{item.keyword.title()}</strong> 
                                    <span style="float: right; color: #28a745;">+{item.weight} pts</span>
                                </div>
                            """, unsafe_allow_html=True)
                        else:
                            st.markdown(f"""
                                <div class="keyword-missing">
                                    ❌ <strong>{item.keyword.title()}</strong> - Not found
                                    <span style="float: right; color: #dc3545;">0/{item.weight} pts</span>
                                </div>
                            """, unsafe_allow_html=True)
            
            # Summary stats
            st.markdown("---")
            found_count = sum(1 for r in analysis["results"] if r.found)
            total_count = len(analysis["results"])
            
            col_a, col_b, col_c = st.columns(3)
            with col_a:
                st.metric("Keywords Found", f"{found_count}/{total_count}")
            with col_b:
                st.metric("Points Earned", f"{analysis['earned_points']}/{analysis['total_possible']}")
            with col_c:
                st.metric("Penalties", f"-{analysis['penalty']}" if analysis['penalty'] > 0 else "None")
        
        elif analyze_button:
            st.warning("Please paste your AI Use Policy text to analyze.")
        else:
            st.info("👈 Paste your policy text and click 'Analyze Policy' to begin.")
            
            # Show what we check for
            st.markdown("---")
            st.subheader("🎯 What We Check For")
            
            checks = [
                ("Liability", "Legal protection clauses"),
                ("Human Review", "Manual oversight processes"),
                ("Data Training", "Training data governance"),
                ("Human-in-the-Loop", "Critical safety requirement"),
                ("Accountability", "Responsibility frameworks"),
                ("Transparency", "Disclosure practices"),
                ("Bias", "Fairness considerations"),
                ("Audit", "Review mechanisms"),
                ("Consent", "User permission protocols"),
                ("Privacy", "Data protection measures"),
            ]
            
            for keyword, description in checks:
                st.markdown(f"• **{keyword}**: {description}")
    
    # Footer
    st.markdown("---")
    st.markdown("""
        <div style="text-align: center; color: #666; padding: 20px;">
            <p>🔒 Cata Risk Lab Policy Auditor | Helping organizations build safer AI practices</p>
        </div>
    """, unsafe_allow_html=True)


if __name__ == "__main__":
    main()