import os
import json
import time
import uuid
import hashlib
from typing import Optional, List, Dict, Any, Tuple

# ---------------------------------------------------------------------
# Core identifiers
# ---------------------------------------------------------------------

ENGINE_NAME = "AI_Sovereign_Sentinel_Core_v1"
AUTHORITY_NAME = "DataClear Sovereign Authority"
SOVEREIGN_VERSION = "1.2-gov-ready"

# Centralised audit log path (JSONL)
AUDIT_LOG_FILE = "sovereign_audit_log.jsonl"


# ---------------------------------------------------------------------
# Fingerprinting
# ---------------------------------------------------------------------


class SovereignFingerprint:
    """
    Minimal fingerprint generator for AI Sovereign Sentinel.
    Creates a JSON file with engine, fingerprint and issued_at.
    """

    def __init__(self, engine: str, output: str = "sovereign_fingerprint.json"):
        self.engine = engine
        self.output = output

    def _build(self) -> Dict[str, Any]:
        issued_at = int(time.time())
        base = f"{self.engine}|{issued_at}"
        fp_hash = hashlib.sha256(base.encode()).hexdigest()
        return {
            "engine": self.engine,
            "fingerprint": fp_hash,
            "issued_at": issued_at,
            "version": "1.0",
        }

    def export(self) -> Tuple[str, str]:
        fp = self._build()
        with open(self.output, "w", encoding="utf-8") as f:
            json.dump(fp, f, indent=4)
        return self.output, fp["fingerprint"]


# ---------------------------------------------------------------------
# Engine registry / validator
# ---------------------------------------------------------------------


class SovereignValidator:
    """
    Simple registry to bind engine name → fingerprint string.
    Stored in sovereign_registry.json
    """

    def __init__(self, path: str = "sovereign_registry.json"):
        self.path = path
        if os.path.exists(self.path):
            with open(self.path, "r", encoding="utf-8") as f:
                self.registry = json.load(f)
        else:
            self.registry = {}

    def register(self, engine: str, fingerprint: str) -> bool:
        self.registry[engine] = {
            "fingerprint": fingerprint,
            "registered_at": int(time.time()),
        }
        with open(self.path, "w", encoding="utf-8") as f:
            json.dump(self.registry, f, indent=4)
        return True

    def validate(self, engine: str, fingerprint: str) -> bool:
        entry = self.registry.get(engine)
        if not entry:
            return False
        return entry.get("fingerprint") == fingerprint


# ---------------------------------------------------------------------
# Attestation
# ---------------------------------------------------------------------


def attestation_signature(att: dict) -> str:
    """
    Unified, stable signature generation for attestation objects.
    """
    base = (
        f"{att['attestation_id']}"
        f"|{att['issued_at']}"
        f"|{att['fingerprint']}"
        f"|{att['issuer']}"
        f"|{att['version']}"
    )
    return hashlib.sha256(base.encode()).hexdigest()


def generate_attestation(fp) -> Dict[str, Any]:
    """
    Create a signed attestation for a given fingerprint (dict or string).
    """
    if isinstance(fp, dict):
        fp_value = fp.get("fingerprint", "")
    else:
        fp_value = fp

    att: Dict[str, Any] = {
        "attestation_id": str(uuid.uuid4()),
        "issued_at": int(time.time()),
        "fingerprint": fp_value,
        "issuer": AUTHORITY_NAME,
        "version": "1.0",
    }

    att["signature"] = attestation_signature(att)
    return att


def verify_attestation(att: dict) -> bool:
    expected = attestation_signature(att)
    return expected == att.get("signature")


# ---------------------------------------------------------------------
# Lineage + Deployment binding (multi-model metadata ready)
# ---------------------------------------------------------------------



def build_lineage_record(
    engine: str,
    fingerprint: str,
    parent_model: Optional[str] = None,
    model_version: str = "1.0",
    data_tags: Optional[List[str]] = None,
    notes: Optional[str] = None,
    provider: Optional[str] = None,
    model_family: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Create a cryptographically-bound lineage record for a model.

    Hash formula (ثابت و قابل بازتولید):
      payload = f"{engine}|{fingerprint}|{model_version}|{parent_model or ''}|{created_at}"
    """

    created_at = int(time.time())

    core_parent = parent_model or ""
    payload = f"{engine}|{fingerprint}|{model_version}|{core_parent}|{created_at}"
    lineage_hash = hashlib.sha256(payload.encode()).hexdigest()

    return {
        "lineage_id": str(uuid.uuid4()),
        "engine": engine,
        "fingerprint": fingerprint,
        "parent_model": parent_model,
        "model_version": model_version,
        "data_tags": data_tags or [],
        "notes": notes,
        "provider": provider,
        "model_family": model_family,
        "created_at": created_at,
        "lineage_hash": lineage_hash,
    }





# ---------------------------------------------------------------------
# Simple risk profile (env + data tags)
# ---------------------------------------------------------------------


def compute_risk_profile(
    data_tags: Optional[List[str]] = None,
    deployment_env: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Very simple, rule-based risk scoring so that certificate
    carries a usable risk_level and risk_score.
    """
    tags = [t.lower() for t in (data_tags or [])]
    env = (deployment_env or "").lower().strip()

    score = 0
    threat_flags: List[str] = []

    # Data-sensitivity
    sensitive_keywords = ["pii", "health", "healthcare", "finance", "bank", "payments"]
    for kw in sensitive_keywords:
        if any(kw in t for t in tags):
            score += 25
            threat_flags.append(f"sensitive_data:{kw}")
            break

    # Internal / confidential
    if any("internal" in t or "confidential" in t for t in tags):
        score += 15
        threat_flags.append("internal_confidential")

    # Environment-based
    if env in ["prod", "production", "live"]:
        score += 40
        threat_flags.append("environment:production")
    elif env in ["staging", "preprod", "test"]:
        score += 20
        threat_flags.append("environment:staging")
    elif env:
        score += 10
        threat_flags.append(f"environment:{env}")

    # Cap score at 100
    score = max(0, min(score, 100))

    if score >= 70:
        level = "HIGH"
    elif score >= 40:
        level = "MEDIUM"
    else:
        level = "LOW"

    return {
        "level": level,
        "score": score,
        "threat_flags": threat_flags,
    }


# ---------------------------------------------------------------------
# Compute-efficiency scoring (approximate)
# ---------------------------------------------------------------------


def estimate_compute_cost(
    model_name: str,
    tokens: int,
    provider: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Tiny heuristic for 'compute efficiency' the government is obsessed with.
    We approximate per-token cost by model family.
    """

    name = (model_name or "").lower()
    provider_norm = (provider or "").lower()

    if any(k in name for k in ["gpt-4", "claude-3-opus", "gemini-1.5-pro"]) or "frontier" in name:
        base_cost = 3.0
        family = "frontier"
    elif any(k in name for k in ["gpt-4o-mini", "gpt-3.5", "sonnet", "llama-3", "mistral"]):
        base_cost = 1.0
        family = "general"
    else:
        base_cost = 0.2
        family = "narrow"

    estimated_units = base_cost * tokens

    if family == "narrow":
        band = "HIGH_EFFICIENCY"
    elif family == "general":
        band = "BALANCED"
    else:
        band = "HIGH_COMPUTE"

    return {
        "provider": provider_norm or None,
        "model_name": model_name,
        "family": family,
        "estimated_compute_units": estimated_units,
        "efficiency_band": band,
        "tokens": tokens,
    }


# ---------------------------------------------------------------------
# Policy Engine (standard & custom)
# ---------------------------------------------------------------------


class PolicyEngine:
    """
    Simple, pluggable policy engine.
    Provides:
      - standard safety templates (LLM safety / restriction)
      - rule-based enforcement
      - support for custom organisation policies
    """

    def __init__(self):
        self.policies: Dict[str, Dict[str, Any]] = {}
        self._load_default_policies()

    def _load_default_policies(self):
        self.policies["SAFE_OUTPUT_BASELINE"] = {
            "id": "SAFE_OUTPUT_BASELINE",
            "description": "Disallow obvious harmful, illegal or highly sensitive content.",
            "rules": {
                "blocked_keywords": [
                    "terrorist",
                    "child sexual",
                    "explosive instructions",
                    "bioweapon",
                    "financial fraud",
                ]
            },
        }
        self.policies["NO_PII_LEAK"] = {
            "id": "NO_PII_LEAK",
            "description": "Reduce risk of leaking obvious PII markers in logs.",
            "rules": {
                "blocked_markers": ["national insurance number", "passport number", "nhs number"]
            },
        }

    def add_policy(self, policy: Dict[str, Any]) -> None:
        pid = policy.get("id")
        if not pid:
            raise ValueError("Policy must have an 'id'")
        self.policies[pid] = policy

    def list_policies(self) -> List[Dict[str, Any]]:
        return list(self.policies.values())

    def evaluate_output(
        self,
        output_text: str,
        active_policy_ids: Optional[List[str]] = None,
    ) -> Dict[str, Any]:
        text = (output_text or "").lower()
        active_policy_ids = active_policy_ids or list(self.policies.keys())

        violations: List[Dict[str, Any]] = []

        for pid in active_policy_ids:
            policy = self.policies.get(pid)
            if not policy:
                continue

            rules = policy.get("rules", {})

            for kw in rules.get("blocked_keywords", []):
                if kw in text:
                    violations.append(
                        {"policy_id": pid, "type": "blocked_keyword", "value": kw}
                    )

            for marker in rules.get("blocked_markers", []):
                if marker in text:
                    violations.append(
                        {"policy_id": pid, "type": "blocked_marker", "value": marker}
                    )

        is_allowed = len(violations) == 0
        return {"is_allowed": is_allowed, "violations": violations}


# ---------------------------------------------------------------------
# Integrity certificate
# ---------------------------------------------------------------------


def generate_integrity_certificate(
    fp,
    att,
    output: str = "sovereign_certificate.json",
    parent_model: Optional[str] = None,
    model_version: str = "1.0",
    data_tags: Optional[List[str]] = None,
    notes: Optional[str] = None,
    deployment_env: Optional[str] = None,
    provider: Optional[str] = None,
    model_family: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Bundle fingerprint + attestation + lineage + risk profile
    into a single integrity certificate file.
    """
    if isinstance(fp, dict):
        fp_value = fp.get("fingerprint", "")
    else:
        fp_value = fp

    lineage = build_lineage_record(
        engine=ENGINE_NAME,
        fingerprint=fp_value,
        parent_model=parent_model,
        model_version=model_version,
        data_tags=data_tags,
        notes=notes,
        provider=provider,
        model_family=model_family,
    )

    risk = compute_risk_profile(data_tags=data_tags, deployment_env=deployment_env)

    issued_at = att["issued_at"]

    cert: Dict[str, Any] = {
        "engine": ENGINE_NAME,
        "sovereign_version": SOVEREIGN_VERSION,
        "fingerprint": fp_value,
        "issued_at": issued_at,
        "attestation": att,
        "lineage": lineage,
        "risk_level": risk["level"],
        "risk_score": risk["score"],
        "threat_flags": risk["threat_flags"],
        "environment": deployment_env or "unspecified",
        "audit_id": str(uuid.uuid4()),
    }

    with open(output, "w", encoding="utf-8") as f:
        json.dump(cert, f, indent=4)

    return cert


# ---------------------------------------------------------------------
# Centralised Audit Log Layer  (JSONL + CSV/JSON export)
# ---------------------------------------------------------------------


def _ensure_audit_log_dir(path: str) -> None:
    directory = os.path.dirname(path)
    if directory and not os.path.exists(directory):
        os.makedirs(directory, exist_ok=True)


def write_audit_event(
    event_type: str,
    engine: str,
    fingerprint: str,
    parent_model: Optional[str],
    model_version: str,
    data_tags: Optional[List[str]],
    risk_level: Optional[str],
    risk_score: Optional[int],
    deployment_env: Optional[str],
    extra: Optional[Dict[str, Any]] = None,
    path: str = AUDIT_LOG_FILE,
) -> Dict[str, Any]:
    """
    Append a single audit event as JSON line into the centralised audit journal.
    Designed to be reusable / exportable for G-Cloud / NCSC type evidence.
    """
    _ensure_audit_log_dir(path)

    entry = {
        "event_id": str(uuid.uuid4()),
        "event_type": event_type,
        "engine": engine,
        "fingerprint": fingerprint,
        "parent_model": parent_model,
        "model_version": model_version,
        "data_tags": data_tags or [],
        "risk_level": risk_level,
        "risk_score": risk_score,
        "deployment_env": deployment_env,
        "timestamp": int(time.time()),
        "extra": extra or {},
    }

    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")

    return entry


def export_audit_log_json(path: str = AUDIT_LOG_FILE) -> List[Dict[str, Any]]:
    """
    Load the audit log as a list of JSON objects.
    """
    if not os.path.exists(path):
        return []

    entries: List[Dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
                entries.append(obj)
            except Exception:
                # skip corrupted lines
                continue
    return entries


def export_audit_log_csv(path: str = AUDIT_LOG_FILE) -> str:
    """
    Export the audit log as a simple CSV string
    (for G-Cloud / NCSC style evidence export).
    """
    entries = export_audit_log_json(path=path)
    if not entries:
        return ""

    headers = list(entries[0].keys())
    lines: List[str] = []

    # header row
    lines.append(",".join(headers))

    # data rows
    for row in entries:
        row_values: List[str] = []
        for h in headers:
            value = row.get(h, "")
            if isinstance(value, list):
                value = ";".join(str(v) for v in value)
            elif isinstance(value, dict):
                value = json.dumps(value)
            row_values.append(str(value))
        lines.append(",".join(row_values))

    return "\n".join(lines)


# ---------------------------------------------------------------------
# Higher-level helper: log model execution with policy & efficiency info
# ---------------------------------------------------------------------


def log_model_execution(
    model_name: str,
    provider: Optional[str],
    deployment_env: str,
    data_tags: Optional[List[str]],
    output_text: str,
    tokens: int,
    parent_model: Optional[str] = None,
    model_version: str = "1.0",
    active_policy_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """
    Convenience helper that:
      - computes risk profile
      - evaluates policy engine
      - estimates compute efficiency
      - writes a single unified audit event
    """
    # risk
    risk = compute_risk_profile(data_tags=data_tags, deployment_env=deployment_env)

    # policy
    pe = PolicyEngine()
    policy_result = pe.evaluate_output(output_text, active_policy_ids=active_policy_ids)

    # efficiency
    efficiency = estimate_compute_cost(
        model_name=model_name,
        tokens=tokens,
        provider=provider,
    )

    extra = {
        "model_name": model_name,
        "provider": provider,
        "policy_result": policy_result,
        "efficiency": efficiency,
    }

    entry = write_audit_event(
        event_type="model_execution",
        engine=ENGINE_NAME,
        fingerprint=efficiency["family"],  # simple placeholder binding
        parent_model=parent_model,
        model_version=model_version,
        data_tags=data_tags,
        risk_level=risk["level"],
        risk_score=risk["score"],
        deployment_env=deployment_env,
        extra=extra,
    )

    return entry


# ------------------------------------------------------------
# Certificate Verification Wrapper (Fix for app.py import)
# ------------------------------------------------------------


def verify_certificate(cert: dict) -> bool:
    """
    Verification wrapper so that app.py can import it.

    A valid certificate requires:
      - attestation to verify successfully
      - fingerprint / engine to be consistent across cert + lineage + attestation
      - lineage_hash to match a recomputed hash from core fields
    """
    try:
        # 1) Attestation must be valid
        att = cert.get("attestation") or {}
        if not verify_attestation(att):
            return False

        # 2) Basic consistency checks
        cert_fp = cert.get("fingerprint")
        cert_engine = cert.get("engine")

        lineage = cert.get("lineage") or {}
        line_fp = lineage.get("fingerprint")
        line_engine = lineage.get("engine")

        if not cert_fp:
            return False

        # fingerprint باید همه‌جا یکی باشد
        if cert_fp != att.get("fingerprint"):
            return False
        if cert_fp != line_fp:
            return False

        # engine اگر در هر دو جا هست باید یکی باشد
        if cert_engine and line_engine and cert_engine != line_engine:
            return False

        # 3) Recompute lineage_hash exactly like build_lineage_record
        engine_val = line_engine or ""
        fp_val = line_fp or ""
        model_ver = (lineage.get("model_version") or "")
        parent_val = (lineage.get("parent_model") or "")
        created_val = lineage.get("created_at") or ""

        payload = f"{engine_val}|{fp_val}|{model_ver}|{parent_val}|{created_val}"
        expected_hash = hashlib.sha256(payload.encode()).hexdigest()

        if lineage.get("lineage_hash") != expected_hash:
            return False

        # اگر به اینجا برسیم، همه چیز ok است
        return True

    except Exception:
        return False