import gradio as gr
from typing import List, Dict, Tuple, Optional
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta
import random
import time
import json
import os

# Mock data generation functions
def generate_threat_data() -> pd.DataFrame:
    """Generate mock threat intelligence data"""
    threats = ["Credential Leak", "Phishing Campaign", "Malware Distribution", "Dark Web Mention", "Vulnerability Exploit"]
    sources = ["Twitter", "Pastebin", "Dark Web Forums", "GitHub", "Public Records"]
    now = datetime.now()
    
    data = []
    for i in range(100):
        days_ago = random.randint(0, 30)
        threat_date = now - timedelta(days=days_ago)
        data.append({
            "threat_type": random.choice(threats),
            "source": random.choice(sources),
            "severity": random.randint(1, 10),
            "confidence": random.randint(50, 100),
            "date": threat_date.strftime("%Y-%m-%d"),
            "affected_entity": f"entity-{random.randint(1, 20)}",
            "status": random.choice(["Active", "Mitigated", "Investigating"])
        })
    return pd.DataFrame(data)

def generate_risk_scores() -> Dict[str, int]:
    """Generate mock risk scores for entities"""
    return {f"entity-{i}": random.randint(1, 100) for i in range(1, 21)}

def generate_alerts() -> List[List]:
    """Generate active alerts"""
    alert_types = ["Credential Exposure", "Suspicious Login", "Data Leak", "Unauthorized Access", "Malware Detected"]
    now = datetime.now()
    alerts = []
    for i in range(5, 0, -1):
        alert_time = now - timedelta(minutes=i*15)
        alerts.append([
            alert_time.strftime("%H:%M"),
            random.choice(alert_types),
            random.randint(1, 5)
        ])
    return alerts

def generate_threat_intel() -> List[Tuple[str, str]]:
    """Generate threat intelligence feed items"""
    intel_items = [
        ("[10:45] New phishing campaign targeting financial sector", "Phishing"),
        ("[09:30] Credential leak detected on paste site", "Credential Leak"),
        ("[08:15] Vulnerability CVE-2023-1234 actively exploited", "Exploit"),
        ("[07:00] Dark web forum discussing company data", "Dark Web"),
        ("[06:30] Suspicious API activity detected", "API Abuse")
    ]
    return intel_items

# Analytics functions
def create_threat_timeline(df: pd.DataFrame) -> go.Figure:
    """Create interactive threat timeline"""
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(text="No data available", xref="paper", yref="paper", showarrow=False)
        return fig
    
    timeline_data = df.groupby(["date", "threat_type"]).size().reset_index(name="count")
    fig = px.line(
        timeline_data,
        x="date",
        y="count",
        color="threat_type",
        title="Threat Activity Timeline",
        labels={"date": "Date", "count": "Threat Count", "threat_type": "Threat Type"}
    )
    fig.update_layout(height=400, hovermode='x unified')
    return fig

def create_risk_heatmap(scores: Dict[str, int]) -> go.Figure:
    """Create risk score heatmap"""
    if not scores:
        fig = go.Figure()
        fig.add_annotation(text="No data available", xref="paper", yref="paper", showarrow=False)
        return fig
    
    df = pd.DataFrame(list(scores.items()), columns=["Entity", "Risk Score"])
    fig = px.imshow(
        [list(scores.values())],
        labels=dict(x="Entities", y="Risk", color="Score"),
        x=list(scores.keys()),
        title="Entity Risk Scores",
        color_continuous_scale="reds"
    )
    fig.update_layout(height=400)
    return fig

def create_threat_distribution(df: pd.DataFrame) -> go.Figure:
    """Create threat type distribution pie chart"""
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(text="No data available", xref="paper", yref="paper", showarrow=False)
        return fig
    
    dist_data = df["threat_type"].value_counts().reset_index()
    dist_data.columns = ["Threat Type", "Count"]
    fig = px.pie(
        dist_data,
        values="Count",
        names="Threat Type",
        title="Threat Type Distribution"
    )
    fig.update_layout(height=350)
    return fig

def create_source_analysis(df: pd.DataFrame) -> go.Figure:
    """Create source analysis bar chart"""
    if df.empty:
        fig = go.Figure()
        fig.add_annotation(text="No data available", xref="paper", yref="paper", showarrow=False)
        return fig
    
    source_data = df["source"].value_counts().reset_index()
    source_data.columns = ["Source", "Count"]
    fig = px.bar(
        source_data,
        x="Source",
        y="Count",
        title="Threats by Source",
        color="Count",
        color_continuous_scale="blues"
    )
    fig.update_layout(height=350)
    return fig

def filter_threat_data(df: pd.DataFrame, threat_types: List[str], days: int) -> pd.DataFrame:
    """Filter threat data based on criteria"""
    if df.empty:
        return df
    
    filtered = df.copy()
    if threat_types:
        filtered = filtered[filtered["threat_type"].isin(threat_types)]
    
    cutoff_date = datetime.now() - timedelta(days=int(days))
    filtered["date"] = pd.to_datetime(filtered["date"])
    filtered = filtered[filtered["date"] >= cutoff_date]
    
    return filtered

def acknowledge_alerts(alerts: List[List]) -> Tuple[List[List], str]:
    """Acknowledge all alerts"""
    acknowledged_count = len(alerts)
    return [], f"✓ Acknowledged {acknowledged_count} alerts"

def export_report(threat_data: pd.DataFrame, risk_scores: Dict[str, int]) -> str:
    """Export report to JSON"""
    report = {
        "generated_at": datetime.now().isoformat(),
        "threat_summary": {
            "total_threats": len(threat_data),
            "threat_types": threat_data["threat_type"].value_counts().to_dict() if not threat_data.empty else {},
            "avg_severity": float(threat_data["severity"].mean()) if not threat_data.empty else 0
        },
        "risk_summary": {
            "entities_monitored": len(risk_scores),
            "high_risk_count": sum(1 for score in risk_scores.values() if score > 70),
            "avg_risk_score": sum(risk_scores.values()) / len(risk_scores) if risk_scores else 0
        }
    }
    
    filename = f"nexusme_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(filename, 'w') as f:
        json.dump(report, f, indent=2)
    
    return f"✓ Report exported to {filename}"

def start_new_investigation() -> str:
    """Start new investigation"""
    investigation_id = f"INV-{datetime.now().strftime('%Y%m%d%H%M%S')}"
    return f"✓ New investigation started: {investigation_id}"

def run_security_scan() -> Tuple[str, str]:
    """Run security scan"""
    scan_id = f"SCAN-{datetime.now().strftime('%Y%m%d%H%M%S')}"
    time.sleep(2)  # Simulate scan
    findings = random.randint(0, 5)
    status = "Complete" if findings == 0 else f"Complete - {findings} findings"
    return scan_id, status

# Core application
with gr.Blocks() as demo:
    # Header with security warning
    gr.Markdown("""
    # 🛡️ Nexusme OSINT Platform
    **Advanced Threat Intelligence Dashboard**  
    *For authorized security personnel only - All access is logged and monitored*
    
    <div style='text-align: center; padding: 10px; background: #fef3c7; border-radius: 5px; margin: 10px 0;'>
    ⚠️ **Security Notice**: This system is for authorized use only. Unauthorized access is prohibited.
    </div>
    """)
    
    # State management
    threat_data_state = gr.State(generate_threat_data())
    risk_scores_state = gr.State(generate_risk_scores())
    alerts_state = gr.State(generate_alerts())
    
    with gr.Tabs():
        # Tab 1: Threat Landscape
        with gr.TabItem("🌐 Threat Landscape", id=1):
            with gr.Row():
                with gr.Column(scale=2):
                    timeline_plot = gr.Plot(label="Threat Activity Timeline")
                    
                    with gr.Row():
                        with gr.Column():
                            threat_dist_plot = gr.Plot(label="Threat Distribution")
                        with gr.Column():
                            source_plot = gr.Plot(label="Source Analysis")
                    
                    refresh_btn = gr.Button("🔄 Refresh Data", variant="primary", size="lg")
                
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("### 🔍 Advanced Filters")
                        threat_type_filter = gr.CheckboxGroup(
                            ["Credential Leak", "Phishing Campaign", "Malware Distribution", "Dark Web Mention", "Vulnerability Exploit"],
                            label="Filter Threat Types",
                            value=["Credential Leak", "Phishing Campaign", "Malware Distribution"]
                        )
                        date_range = gr.Slider(1, 30, 7, label="Days to Analyze", step=1)
                        apply_filter_btn = gr.Button("Apply Filters", variant="secondary")
                    
                    with gr.Group():
                        gr.Markdown("### 📊 Statistics")
                        stats_json = gr.JSON(label="Threat Statistics")
        
        # Tab 2: Risk Assessment
        with gr.TabItem("⚠️ Risk Assessment", id=2):
            with gr.Row():
                with gr.Column(scale=2):
                    heatmap = gr.Plot(label="Entity Risk Scores")
                    risk_refresh = gr.Button("🔄 Refresh Scores", variant="primary")
                
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("### 📈 Risk Summary")
                        risk_summary = gr.JSON(label="Risk Metrics")
                    
                    with gr.Group():
                        gr.Markdown("### 🎯 High Risk Entities")
                        high_risk_table = gr.Dataframe(
                            headers=["Entity", "Risk Score", "Status"],
                            datatype=["str", "number", "str"],
                            interactive=False,
                            label="High Risk (>70)"
                        )
        
        # Tab 3: Alerts & Intelligence
        with gr.TabItem("🚨 Alerts & Intel", id=3):
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("### 🚨 Active Alerts")
                        alert_table = gr.Dataframe(
                            headers=["Time", "Alert Type", "Severity"],
                            datatype=["str", "str", "number"],
                            interactive=False,
                            label="Current Alerts"
                        )
                        acknowledge_btn = gr.Button("✓ Acknowledge All", variant="secondary")
                        ack_status = gr.Textbox(label="Status", interactive=False)
                    
                    with gr.Group():
                        gr.Markdown("### 📡 Threat Intelligence Feed")
                        intel_feed = gr.HighlightedText(
                            value=generate_threat_intel(),
                            color_map={
                                "Phishing": "#FF6B6B",
                                "Credential Leak": "#FFA500",
                                "Exploit": "#8B0000",
                                "Dark Web": "#4B0082",
                                "API Abuse": "#DC143C"
                            },
                            show_legend=True,
                            label=""
                        )
                
                with gr.Column(scale=1):
                    with gr.Group():
                        gr.Markdown("### ⚡ Quick Actions")
                        with gr.Row():
                            export_btn = gr.Button("📄 Export Report", variant="secondary")
                            new_search = gr.Button("🔍 New Investigation", variant="secondary")
                        scan_now = gr.Button("🔒 Scan Now", variant="primary", size="lg")
                        
                        export_status = gr.Textbox(label="Export Status", interactive=False)
                        investigation_status = gr.Textbox(label="Investigation Status", interactive=False)
                        scan_id_output = gr.Textbox(label="Scan ID", interactive=False)
                        scan_status_output = gr.Textbox(label="Scan Status", interactive=False)
    
    # Event handlers
    def update_all_visualizations(threat_data, risk_scores):
        """Update all visualization components"""
        timeline = create_threat_timeline(threat_data)
        threat_dist = create_threat_distribution(threat_data)
        source_analysis = create_source_analysis(threat_data)
        heatmap_viz = create_risk_heatmap(risk_scores)
        
        stats = {
            "total_threats": len(threat_data),
            "unique_entities": threat_data["affected_entity"].nunique() if not threat_data.empty else 0,
            "avg_severity": float(threat_data["severity"].mean()) if not threat_data.empty else 0,
            "avg_confidence": float(threat_data["confidence"].mean()) if not threat_data.empty else 0
        }
        
        risk_metrics = {
            "entities_monitored": len(risk_scores),
            "high_risk_count": sum(1 for score in risk_scores.values() if score > 70),
            "avg_risk_score": round(sum(risk_scores.values()) / len(risk_scores), 2) if risk_scores else 0,
            "max_risk_score": max(risk_scores.values()) if risk_scores else 0
        }
        
        high_risk_entities = [
            [entity, score, "Critical" if score > 90 else "High"]
            for entity, score in risk_scores.items()
            if score > 70
        ][:10]
        
        return [
            timeline, threat_dist, source_analysis, heatmap_viz,
            stats, risk_metrics, high_risk_entities
        ]
    
    # Initial load
    demo.load(
        fn=update_all_visualizations,
        inputs=[threat_data_state, risk_scores_state],
        outputs=[
            timeline_plot, threat_dist_plot, source_plot, heatmap,
            stats_json, risk_summary, high_risk_table
        ]
    )
    
    # Refresh threat data
    refresh_btn.click(
        fn=generate_threat_data,
        outputs=threat_data_state
    ).then(
        fn=update_all_visualizations,
        inputs=[threat_data_state, risk_scores_state],
        outputs=[
            timeline_plot, threat_dist_plot, source_plot, heatmap,
            stats_json, risk_summary, high_risk_table
        ]
    )
    
    # Apply filters
    apply_filter_btn.click(
        fn=filter_threat_data,
        inputs=[threat_data_state, threat_type_filter, date_range],
        outputs=threat_data_state
    ).then(
        fn=lambda td, rs: [
            create_threat_timeline(td),
            create_threat_distribution(td),
            create_source_analysis(td),
            create_risk_heatmap(rs),
            {
                "total_threats": len(td),
                "unique_entities": td["affected_entity"].nunique() if not td.empty else 0,
                "avg_severity": float(td["severity"].mean()) if not td.empty else 0
            },
            risk_summary.value,
            high_risk_table.value
        ],
        inputs=[threat_data_state, risk_scores_state],
        outputs=[
            timeline_plot, threat_dist_plot, source_plot, heatmap,
            stats_json, risk_summary, high_risk_table
        ]
    )
    
    # Refresh risk scores
    risk_refresh.click(
        fn=generate_risk_scores,
        outputs=risk_scores_state
    ).then(
        fn=update_all_visualizations,
        inputs=[threat_data_state, risk_scores_state],
        outputs=[
            timeline_plot, threat_dist_plot, source_plot, heatmap,
            stats_json, risk_summary, high_risk_table
        ]
    )
    
    # Acknowledge alerts
    acknowledge_btn.click(
        fn=acknowledge_alerts,
        inputs=alerts_state,
        outputs=[alerts_state, ack_status]
    ).then(
        fn=lambda: [],
        outputs=alert_table
    )
    
    # Export report
    export_btn.click(
        fn=export_report,
        inputs=[threat_data_state, risk_scores_state],
        outputs=export_status
    )
    
    # New investigation
    new_search.click(
        fn=start_new_investigation,
        outputs=investigation_status
    )
    
    # Security scan
    scan_now.click(
        fn=run_security_scan,
        outputs=[scan_id_output, scan_status_output]
    )
    
    # Footer with compliance info
    gr.Markdown("""
    <div style='text-align: center; font-size: 0.8em; margin-top: 20px; padding: 10px; background: #f3f4f6; border-radius: 5px;'>
    <strong>Nexusme OSINT Platform v1.0</strong> | 
    <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank' style='color: #4f46e5; text-decoration: none;'>Built with anycoder</a> | 
    All activities logged and monitored for security compliance
    </div>
    """)

# Launch with security-focused settings
if __name__ == "__main__":
    demo.launch(
        auth=("admin", "securepassword123"),
        auth_message="Please authenticate with your Nexusme credentials",
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=False,
        footer_links=[
            {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
            {"label": "Privacy Policy", "url": "#"},
            {"label": "Terms of Service", "url": "#"},
            {"label": "Documentation", "url": "#"}
        ],
        theme=gr.themes.Soft(
            primary_hue="indigo",
            secondary_hue="blue",
            neutral_hue="slate",
            font=[gr.themes.GoogleFont("Roboto"), "ui-sans-serif", "system-ui"],
            text_size="md",
            spacing_size="md",
            radius_size="md"
        ).set(
            button_primary_background_fill="*primary_600",
            button_primary_background_fill_hover="*primary_700",
            block_title_text_weight="600",
        )
    )