# hf_demo.py – ARF v4 dashboard for Hugging Face Spaces
import logging
import uuid
from datetime import datetime, timezone

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
import gradio as gr

# ARF v4 imports
from agentic_reliability_framework.core.governance.risk_engine import RiskEngine
from agentic_reliability_framework.runtime.memory import create_faiss_index, RAGGraphMemory
from agentic_reliability_framework.runtime.memory.constants import MemoryConstants

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="ARF v4 API with Memory")

# Enable CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["https://arf-frontend-sandy.vercel.app"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------------------------------------------------------------------------
# Initialize ARF components
# ---------------------------------------------------------------------------

risk_engine = RiskEngine()

faiss_index = create_faiss_index(dim=MemoryConstants.VECTOR_DIM)
memory = RAGGraphMemory(faiss_index)

# ---------------------------------------------------------------------------
# Decision Memory (NEW)
# ---------------------------------------------------------------------------

decision_history = []
risk_history = []

# ---------------------------------------------------------------------------
# API Endpoints
# ---------------------------------------------------------------------------

@app.get("/")
async def root():
    return {
        "service": "ARF OSS API",
        "version": "4.0.0",
        "status": "operational",
        "memory_stats": memory.get_graph_stats() if memory.has_historical_data() else "empty",
    }

@app.get("/health")
async def health():
    return {"status": "ok", "version": "4.0.0"}

# ---------------------------------------------------------------------------
# Demo Intent Object
# ---------------------------------------------------------------------------

class _DemoIntent:
    environment = "dev"
    deployment_target = "dev"
    service_name = "demo"

# ---------------------------------------------------------------------------
# Risk Calculation
# ---------------------------------------------------------------------------

def _calculate_demo_risk():

    intent = _DemoIntent()

    risk_value, explanation, contributions = risk_engine.calculate_risk(
        intent=intent,
        cost_estimate=None,
        policy_violations=[],
    )

    decision = "approve"
    if risk_value > 0.8:
        decision = "deny"
    elif risk_value > 0.2:
        decision = "escalate"

    decision_id = str(uuid.uuid4())

    decision_record = {
        "id": decision_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "risk": float(risk_value),
        "decision": decision
    }

    decision_history.append(decision_record)

    risk_history.append(
        {
            "timestamp": decision_record["timestamp"],
            "risk": float(risk_value)
        }
    )

    return {
        "risk": float(risk_value),
        "status": "critical" if risk_value > 0.8 else "normal",
        "explanation": explanation,
        "contributions": contributions,
        "decision_id": decision_id,
        "decision": decision
    }

# ---------------------------------------------------------------------------
# Outcome Feedback
# ---------------------------------------------------------------------------

def record_outcome(success: bool):

    if not decision_history:
        return {"error": "no decisions yet"}

    last_decision = decision_history[-1]

    intent = _DemoIntent()

    try:
        risk_engine.update_outcome(intent, success)
    except Exception as e:
        logger.exception("Outcome update failed")

    last_decision["outcome"] = "success" if success else "failure"

    return {
        "decision_id": last_decision["id"],
        "outcome": last_decision["outcome"],
        "timestamp": datetime.now(timezone.utc).isoformat()
    }

# ---------------------------------------------------------------------------
# Snapshot functions
# ---------------------------------------------------------------------------

def get_risk_snapshot():

    try:
        snapshot = _calculate_demo_risk()

        snapshot["timestamp"] = datetime.now(timezone.utc).isoformat()

        return snapshot

    except Exception as e:

        logger.exception("Failed to compute risk snapshot")

        return {
            "status": "error",
            "error": str(e),
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

def get_health_snapshot():

    return {
        "status": "ok",
        "version": "4.0.0",
        "service": "ARF OSS API",
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

def get_memory_snapshot():

    try:

        if memory.has_historical_data():

            stats = memory.get_graph_stats()

            return {
                "status": "ok",
                "memory_stats": stats,
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }

        return {
            "status": "empty",
            "memory_stats": "No historical memory yet.",
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

    except Exception as e:

        logger.exception("Failed to compute memory snapshot")

        return {
            "status": "error",
            "error": str(e),
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }

def get_decision_history():
    return decision_history[-10:]

# ---------------------------------------------------------------------------
# Gradio Dashboard
# ---------------------------------------------------------------------------

with gr.Blocks(title="ARF v4 Demo") as demo:

    gr.Markdown("# Agentic Reliability Framework v4")
    gr.Markdown("### Probabilistic Infrastructure Governance")

    with gr.Row():
        health_output = gr.JSON(label="Health")
        risk_output = gr.JSON(label="Current Risk")

    with gr.Row():
        memory_output = gr.JSON(label="Memory Stats")

    with gr.Row():
        decision_output = gr.JSON(label="Recent Decisions")

    with gr.Row():
        refresh_btn = gr.Button("Evaluate Intent")
        success_btn = gr.Button("Action Succeeded")
        fail_btn = gr.Button("Action Failed")

    refresh_btn.click(fn=get_risk_snapshot, outputs=risk_output)
    success_btn.click(fn=lambda: record_outcome(True), outputs=decision_output)
    fail_btn.click(fn=lambda: record_outcome(False), outputs=decision_output)

    with gr.Row():
        health_btn = gr.Button("Refresh Health")
        memory_btn = gr.Button("Refresh Memory")
        history_btn = gr.Button("Show Decision History")

    health_btn.click(fn=get_health_snapshot, outputs=health_output)
    memory_btn.click(fn=get_memory_snapshot, outputs=memory_output)
    history_btn.click(fn=get_decision_history, outputs=decision_output)

    demo.load(fn=get_health_snapshot, outputs=health_output)
    demo.load(fn=get_risk_snapshot, outputs=risk_output)
    demo.load(fn=get_memory_snapshot, outputs=memory_output)

# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")