"""
eaigins_sdk.risk
────────────────
Stage 4: Risk Scoring Engine

Implements the weighted formula from the paper (Section 4.3):

    Risk Score = (DataSensitivity × W1) + (FinancialImpact × W2)
               + (RegulatoryExposure × W3) + (PolicyViolationProb × W4)

Default weights: 0.30 / 0.25 / 0.25 / 0.20  (sum = 1.0)

Decision thresholds:
    LOW      [0.00, 0.35) → ALLOW
    MEDIUM   [0.35, 0.65) → ESCALATE
    HIGH     [0.65, 0.85) → BLOCK
    CRITICAL [0.85, 1.00] → BLOCK
"""

from __future__ import annotations
from typing import Optional
from .models import (
    AgentAction, FirewallResult, PolicyResult,
    RiskScore, RiskLevel, ExecutionDecision,
)

# ── Default weights ───────────────────────────────────────────────────────────
DEFAULT_WEIGHTS: dict[str, float] = {
    "data_sensitivity":      0.30,
    "financial_impact":      0.25,
    "regulatory_exposure":   0.25,
    "policy_violation_prob": 0.20,
}

DEFAULT_THRESHOLDS: dict[RiskLevel, tuple[float, float]] = {
    RiskLevel.LOW:      (0.00, 0.35),
    RiskLevel.MEDIUM:   (0.35, 0.65),
    RiskLevel.HIGH:     (0.65, 0.85),
    RiskLevel.CRITICAL: (0.85, 1.01),
}

# ── Built-in data sensitivity scores ─────────────────────────────────────────
DATA_SENSITIVITY: dict[str, float] = {
    # Financial / AML
    "transaction_history": 0.80, "customer_full_profile": 0.90,
    "banking_records": 0.85, "kyc_documents": 0.80,
    "account_balances": 0.75, "wire_transfer_details": 0.90,
    "payment_records": 0.85, "fraud_patterns": 0.50,
    "aml_alerts": 0.55, "investigation_notes": 0.60,
    "compliance_reports": 0.50, "regulatory_guidelines": 0.10,
    "public_sanctions_list": 0.15, "fatf_typologies": 0.10,
    "internal_faq": 0.05,
    # Customer care
    "crm_customer_profile": 0.60, "purchase_history": 0.40,
    "order_management_db": 0.20, "shipping_tracker": 0.10,
    "return_policy_db": 0.05, "loyalty_program_db": 0.20,
    "support_ticket_db": 0.30,
    # Healthcare
    "patient_ehr": 0.95, "patient_vitals": 0.75,
    "lab_results": 0.85, "imaging_results": 0.85,
    "psychiatric_records": 0.98, "pharmacy_inventory": 0.50,
    "drug_interaction_db": 0.30, "clinical_guidelines_db": 0.05,
    "public_medical_literature": 0.02, "clinical_history": 0.90,
}

TOOL_FINANCIAL_IMPACT: dict[str, float] = {
    # Financial
    "initiate_wire_transfer": 0.95, "freeze_account": 0.80,
    "close_account": 0.85, "block_customer": 0.75,
    "file_sar": 0.70, "submit_regulatory_report": 0.65,
    "flag_transaction": 0.30, "generate_investigation_report": 0.20,
    "get_transaction_history": 0.10, "get_customer_full_profile": 0.30,
    "query_sanctions_db": 0.10, "add_investigation_note": 0.10,
    "send_to_external_api": 0.90,
    # Healthcare
    "dispense_medication": 0.95, "update_medication_order": 0.90,
    "order_diagnostic_test": 0.40, "generate_prescription_recommendation": 0.50,
    "generate_treatment_recommendation": 0.55, "export_patient_data": 0.85,
    "get_patient_full_record": 0.30, "add_chart_note": 0.10,
    "query_clinical_guidelines": 0.05,
    # Customer care
    "process_refund": 0.55, "apply_account_credits": 0.40,
    "apply_bulk_discount": 0.90, "close_customer_account": 0.60,
    "delete_customer_account": 0.70, "update_payment_method": 0.75,
    "export_customer_data": 0.85, "get_customer_profile": 0.20,
    "get_order_status": 0.05, "create_support_ticket": 0.05,
}

TOOL_REGULATORY_EXPOSURE: dict[str, float] = {
    # Financial
    "file_sar": 0.90, "submit_regulatory_report": 0.85,
    "initiate_wire_transfer": 0.80, "freeze_account": 0.70,
    "close_account": 0.70, "send_to_external_api": 0.75,
    "get_customer_full_profile": 0.60, "generate_investigation_report": 0.40,
    "flag_transaction": 0.20, "get_transaction_history": 0.25,
    "query_sanctions_db": 0.15, "add_investigation_note": 0.10,
    # Healthcare
    "dispense_medication": 0.95, "update_medication_order": 0.95,
    "export_patient_data": 0.90, "order_diagnostic_test": 0.55,
    "generate_prescription_recommendation": 0.60,
    "generate_treatment_recommendation": 0.65,
    "get_patient_full_record": 0.70, "add_chart_note": 0.20,
    "query_clinical_guidelines": 0.05,
    # Customer care
    "delete_customer_account": 0.80, "export_customer_data": 0.85,
    "update_payment_method": 0.80, "apply_bulk_discount": 0.40,
    "process_refund": 0.30, "close_customer_account": 0.50,
    "apply_account_credits": 0.20, "get_customer_profile": 0.35,
    "get_order_status": 0.05, "create_support_ticket": 0.05,
}


class RiskEngine:
    """
    Weighted risk scoring engine implementing the formula from the paper.

    Usage::

        from eaigins_sdk.risk import RiskEngine

        engine = RiskEngine()
        risk   = engine.score(action, firewall_result, policy_result)
        verdict = engine.get_decision(risk)   # ALLOW / ESCALATE / BLOCK

    Custom weights (must sum to 1.0)::

        engine = RiskEngine(weights={
            "data_sensitivity":      0.40,
            "financial_impact":      0.20,
            "regulatory_exposure":   0.20,
            "policy_violation_prob": 0.20,
        })

    Extend built-in lookup maps::

        engine = RiskEngine(
            extra_data_sensitivity={"my_proprietary_db": 0.70},
            extra_financial_impact={"my_payment_tool":   0.80},
        )
    """

    def __init__(
        self,
        weights:                   Optional[dict[str, float]] = None,
        thresholds:                Optional[dict[RiskLevel, tuple[float, float]]] = None,
        extra_data_sensitivity:    Optional[dict[str, float]] = None,
        extra_financial_impact:    Optional[dict[str, float]] = None,
        extra_regulatory_exposure: Optional[dict[str, float]] = None,
    ):
        self.weights    = weights    or DEFAULT_WEIGHTS
        self.thresholds = thresholds or DEFAULT_THRESHOLDS
        total = sum(self.weights.values())
        if abs(total - 1.0) > 0.001:
            raise ValueError(f"Risk weights must sum to 1.0, got {total:.4f}")

        self._ds_map  = {**DATA_SENSITIVITY,         **(extra_data_sensitivity    or {})}
        self._fi_map  = {**TOOL_FINANCIAL_IMPACT,    **(extra_financial_impact    or {})}
        self._re_map  = {**TOOL_REGULATORY_EXPOSURE, **(extra_regulatory_exposure or {})}

    # ── Public API ────────────────────────────────────────────────────────────

    def score(
        self,
        action:          AgentAction,
        firewall_result: FirewallResult,
        policy_result:   PolicyResult,
    ) -> RiskScore:
        """Compute the weighted risk score for ``action``."""
        ds  = self._data_sensitivity(action)
        fi  = self._financial_impact(action)
        re_ = self._regulatory_exposure(action)
        pv  = self._policy_violation(firewall_result, policy_result)

        total = (
            ds  * self.weights["data_sensitivity"]      +
            fi  * self.weights["financial_impact"]      +
            re_ * self.weights["regulatory_exposure"]   +
            pv  * self.weights["policy_violation_prob"]
        )
        total = min(max(total, 0.0), 1.0)
        level = self._classify(total)

        return RiskScore(
            total_score               = round(total, 4),
            data_sensitivity_score    = round(ds,  4),
            financial_impact_score    = round(fi,  4),
            regulatory_exposure_score = round(re_, 4),
            policy_violation_score    = round(pv,  4),
            risk_level                = level,
            breakdown={
                "data_sensitivity":    {"raw": ds,  "weighted": round(ds  * self.weights["data_sensitivity"],      4)},
                "financial_impact":    {"raw": fi,  "weighted": round(fi  * self.weights["financial_impact"],      4)},
                "regulatory_exposure": {"raw": re_, "weighted": round(re_ * self.weights["regulatory_exposure"],   4)},
                "policy_violation":    {"raw": pv,  "weighted": round(pv  * self.weights["policy_violation_prob"], 4)},
            },
        )

    def get_decision(self, risk: RiskScore) -> ExecutionDecision:
        """Map a risk score to ALLOW / ESCALATE / BLOCK."""
        if risk.risk_level in (RiskLevel.HIGH, RiskLevel.CRITICAL):
            return ExecutionDecision.BLOCK
        if risk.risk_level == RiskLevel.MEDIUM:
            return ExecutionDecision.ESCALATE
        return ExecutionDecision.ALLOW

    # ── Private factor scorers ────────────────────────────────────────────────

    def _data_sensitivity(self, action: AgentAction) -> float:
        scores = []
        for src in action.context_sources:
            key = src.lower().replace(" ", "_").replace("-", "_")
            if key in self._ds_map:
                scores.append(self._ds_map[key])
            else:
                scores.append(next((v for k, v in self._ds_map.items()
                                    if k in key or key in k), 0.3))
        tool_key = action.tool_name.lower()
        for k, v in self._ds_map.items():
            if k in tool_key:
                scores.append(v)
                break
        return max(scores) if scores else 0.1

    def _financial_impact(self, action: AgentAction) -> float:
        base = self._fi_map.get(action.tool_name, 0.1)
        amt  = action.tool_params.get("amount", 0)
        if   amt > 1_000_000: base = min(base + 0.3, 1.0)
        elif amt > 100_000:   base = min(base + 0.2, 1.0)
        elif amt > 10_000:    base = min(base + 0.1, 1.0)
        if action.tool_params.get("patient_count",  0) > 100:  base = min(base + 0.2, 1.0)
        if action.tool_params.get("customer_count", 0) > 1000: base = min(base + 0.2, 1.0)
        return base

    def _regulatory_exposure(self, action: AgentAction) -> float:
        return self._re_map.get(action.tool_name, 0.1)

    def _policy_violation(self, fw: FirewallResult, pol: PolicyResult) -> float:
        score  = fw.threat_score * 0.4
        score += min(len(pol.violated_policies)  * 0.20, 0.50)
        score += min(len(pol.required_approvals) * 0.15, 0.30)
        return min(score, 1.0)

    def _classify(self, score: float) -> RiskLevel:
        for level, (lo, hi) in self.thresholds.items():
            if lo <= score < hi:
                return level
        return RiskLevel.CRITICAL
