from __future__ import annotations

import hashlib
import json
import sqlite3
import time
from dataclasses import dataclass, asdict
from typing import Any, Dict, List, Optional, cast

from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives import serialization


def _hash(data: Any) -> str:
    h = hashlib.sha256()
    h.update(json.dumps(data, sort_keys=True, default=str).encode())
    return h.hexdigest()


def _canonical_payload(data: Any) -> bytes:
    return json.dumps(data, sort_keys=True, default=str, separators=(",", ":")).encode()


@dataclass
class LocalProblem:
    problem_id: str
    assets: List[str]
    objectives: List[str]


@dataclass
class SharedVariable:
    name: str
    value: Any
    version: int = 0


@dataclass
class Attestation:
    signer_pubkey_pem: bytes
    signature: bytes

    def to_dict(self) -> Dict[str, Any]:
        return {
            "signer_pubkey_pem": self.signer_pubkey_pem.decode(),
            "signature": self.signature.hex(),
        }

    def verify(self, message: bytes) -> bool:
        public_key = serialization.load_pem_public_key(self.signer_pubkey_pem)
        try:
            cast(Any, public_key).verify(self.signature, message)
            return True
        except Exception:
            return False


@dataclass
class PlanDelta:
    delta_id: str
    changes: Dict[str, Any]
    timestamp: float
    delta_hash: str = ""
    attestation: Optional[Attestation] = None

    def compute_hash(self) -> str:
        self.delta_hash = _hash({"delta_id": self.delta_id, "changes": self.changes, "timestamp": self.timestamp})
        return self.delta_hash

    def to_dict(self) -> Dict[str, Any]:
        payload = asdict(self)
        if self.attestation is not None:
            payload["attestation"] = self.attestation.to_dict()
        return payload


@dataclass
class ProvenanceEvent:
    event_id: str
    experiment_id: str
    actor: str
    timestamp: float
    contract_id: str
    delta_hash: str
    attestation: Attestation
    policy_flags: List[str]

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["attestation"] = self.attestation.to_dict()
        return d


@dataclass
class ProvenanceRecord:
    event: ProvenanceEvent
    verified: bool


class PortaLedger:
    def __init__(self, db_path: str = ":memory:"):
        self.conn = sqlite3.connect(db_path)
        self._init_schema()
        self._signer = Ed25519PrivateKey.generate()

    def _init_schema(self):
        cur = self.conn.cursor()
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS provenance_events (
              event_id TEXT PRIMARY KEY,
              experiment_id TEXT,
              actor TEXT,
              timestamp REAL,
              contract_id TEXT,
              delta_hash TEXT,
              signer_pubkey_pem BLOB,
              attestation BLOB,
              policy_flags TEXT
            )
            """
        )
        cur.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_provenance_experiment
            ON provenance_events (experiment_id, timestamp)
            """
        )
        cur.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_provenance_contract
            ON provenance_events (contract_id, timestamp)
            """
        )
        self.conn.commit()

    def sign_delta(self, delta: PlanDelta) -> PlanDelta:
        delta.compute_hash()
        # Sign the delta hash as a simple attestation
        data = delta.delta_hash.encode()
        signature = self._signer.sign(data)
        pub = self._signer.public_key().public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        delta.attestation = Attestation(signer_pubkey_pem=pub, signature=signature)
        return delta

    def verify_delta(self, delta: PlanDelta) -> bool:
        if not delta.delta_hash:
            delta.compute_hash()
        if not delta.attestation:
            return False
        return delta.attestation.verify(delta.delta_hash.encode())

    def log_provenance(self, event: ProvenanceEvent) -> None:
        cur = self.conn.cursor()
        cur.execute(
            "INSERT OR REPLACE INTO provenance_events (event_id, experiment_id, actor, timestamp, contract_id, delta_hash, signer_pubkey_pem, attestation, policy_flags) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                event.event_id,
                event.experiment_id,
                event.actor,
                event.timestamp,
                event.contract_id,
                event.delta_hash,
                event.attestation.signer_pubkey_pem if event.attestation else None,
                event.attestation.signature if event.attestation else None,
                json.dumps(event.policy_flags),
            ),
        )
        self.conn.commit()

    def record_event(self, experiment_id: str, actor: str, contract_id: str, delta: PlanDelta, policy_flags: Optional[List[str]] = None) -> ProvenanceEvent:
        signed = self.sign_delta(delta)
        attestation = signed.attestation
        assert attestation is not None
        event = ProvenanceEvent(
            event_id=_hash({"experiment_id": experiment_id, "contract_id": contract_id, "delta_hash": signed.delta_hash, "ts": time.time()}),
            experiment_id=experiment_id,
            actor=actor,
            timestamp=time.time(),
            contract_id=contract_id,
            delta_hash=signed.delta_hash,
            attestation=attestation,
            policy_flags=policy_flags or [],
        )
        self.log_provenance(event)
        return event

    def fetch_event(self, event_id: str) -> Optional[ProvenanceEvent]:
        cur = self.conn.cursor()
        cur.execute("SELECT event_id, experiment_id, actor, timestamp, contract_id, delta_hash, signer_pubkey_pem, attestation, policy_flags FROM provenance_events WHERE event_id = ?", (event_id,))
        row = cur.fetchone()
        if not row:
            return None
        attestation = Attestation(signer_pubkey_pem=row[6] or b"", signature=row[7] or b"")
        event = ProvenanceEvent(event_id=row[0], experiment_id=row[1], actor=row[2], timestamp=row[3], contract_id=row[4], delta_hash=row[5], attestation=attestation, policy_flags=json.loads(row[8]))
        return event

    def query_events(self, experiment_id: Optional[str] = None, actor: Optional[str] = None, contract_id: Optional[str] = None) -> List[ProvenanceEvent]:
        clauses = []
        params: List[Any] = []
        if experiment_id is not None:
            clauses.append("experiment_id = ?")
            params.append(experiment_id)
        if actor is not None:
            clauses.append("actor = ?")
            params.append(actor)
        if contract_id is not None:
            clauses.append("contract_id = ?")
            params.append(contract_id)
        query = "SELECT event_id, experiment_id, actor, timestamp, contract_id, delta_hash, signer_pubkey_pem, attestation, policy_flags FROM provenance_events"
        if clauses:
            query += " WHERE " + " AND ".join(clauses)
        query += " ORDER BY timestamp ASC"
        cur = self.conn.cursor()
        cur.execute(query, tuple(params))
        rows = cur.fetchall()
        events: List[ProvenanceEvent] = []
        for row in rows:
            events.append(
                ProvenanceEvent(
                    event_id=row[0],
                    experiment_id=row[1],
                    actor=row[2],
                    timestamp=row[3],
                    contract_id=row[4],
                    delta_hash=row[5],
                    attestation=Attestation(signer_pubkey_pem=row[6] or b"", signature=row[7] or b""),
                    policy_flags=json.loads(row[8]),
                )
            )
        return events

    def verify_event(self, event_id: str) -> bool:
        event = self.fetch_event(event_id)
        if event is None:
            return False
        if not event.attestation.signer_pubkey_pem or not event.attestation.signature:
            return False
        return event.attestation.verify(event.delta_hash.encode())

    def aggregate_privacy_budget(self, events: List[ProvenanceEvent]) -> Dict[str, Any]:
        return {
            "event_count": len(events),
            "unique_actors": sorted({event.actor for event in events}),
            "delta_hashes": [event.delta_hash for event in events],
        }