#!/usr/bin/env python3
"""
ProverEngine — orchestrates the generate-compile-feedback loop.

Ties together:
  - DeepSeekProver (or any Backend with a .generate() method)
  - The language-proof-server for remote compilation
  - The pre-filter for trivial/non-trivial classification
  - Receipt emission for audit trail

Usage:
    from lean_proof import DeepSeekProver, ProverEngine

    prover = DeepSeekProver(backend="ollama")
    engine = ProverEngine(prover)
    result = engine.prove(
        theorem_statement="theorem sq_eq (x : ℕ) : x * x = x ^ 2 := by",
        context="import Mathlib\nopen Nat",
        lake_workdir="/path/to/semantics",
    )
    print(f"Proved: {result.passed}, iterations: {result.iteration}")
"""

from __future__ import annotations

import hashlib
import json
import os
import time
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

from .deepseek_prover import CandidateResult, DeepSeekProver


@dataclass
class ProofResult:
    theorem: str
    passed: bool
    iterations: int
    code: str = ""
    compile_log: str = ""
    model: str = ""
    latency_ms: float = 0.0
    receipt_sha256: str = ""

    def to_receipt(self) -> dict:
        return {
            "schema": "lean_proof_receipt_v1",
            "theorem": self.theorem,
            "passed": self.passed,
            "iterations": self.iterations,
            "model": self.model,
            "latency_ms": self.latency_ms,
            "code_hash": self.receipt_sha256,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }


class ProverEngine:
    """Orchestrates the generate-compile-feedback loop."""

    def __init__(
        self,
        prover: DeepSeekProver,
        log_dir: Optional[Path] = None,
        receipt_dir: Optional[Path] = None,
    ):
        self.prover = prover
        self.log_dir = log_dir or Path("/tmp/lean_proof_logs")
        self.log_dir.mkdir(parents=True, exist_ok=True)
        self.receipt_dir = receipt_dir or (
            Path(__file__).resolve().parents[3]
            / "shared-data"
            / "artifacts"
            / "deepseek_prover"
        )
        self.receipt_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------

    def prove(
        self,
        theorem_statement: str,
        context: str = "",
        max_iterations: int = 5,
        lean_file: str = "/tmp/proof_target.lean",
        lake_workdir: Optional[str] = None,
        emit_receipt: bool = True,
    ) -> ProofResult:
        """Single-theorem proof search loop.

        Args:
            theorem_statement: The Lean theorem to prove.
            context: Surrounding module code.
            max_iterations: Generate-compile cycles.
            lean_file: Temp file for compilation.
            lake_workdir: Working directory for lake build.
            emit_receipt: Write a JSON receipt on completion.

        Returns:
            ProofResult with pass/fail status.
        """
        result = self.prover.prove(
            theorem_statement=theorem_statement,
            context=context,
            max_iterations=max_iterations,
            lean_file=lean_file,
            lake_workdir=lake_workdir,
        )

        proof_result = ProofResult(
            theorem=theorem_statement.split("\n")[0][:120],
            passed=result.passed,
            iterations=result.iteration,
            code=result.code,
            compile_log=result.compile_log,
            model=result.model,
            latency_ms=result.latency_ms,
            receipt_sha256=hashlib.sha256(result.code.encode()).hexdigest(),
        )

        if emit_receipt:
            self._write_receipt(proof_result)

        return proof_result

    # ------------------------------------------------------------------
    # Batch mode
    # ------------------------------------------------------------------

    def prove_batch(
        self,
        theorems: list[tuple[str, str]],
        context: str = "",
        lake_workdir: Optional[str] = None,
    ) -> list[ProofResult]:
        """Prove multiple theorems in sequence.

        Args:
            theorems: List of (theorem_name, theorem_statement) pairs.
            context: Shared context for all theorems.
            lake_workdir: Working directory for compilation.

        Returns:
            List of ProofResult objects.
        """
        results = []
        for name, statement in theorems:
            r = self.prove(
                theorem_statement=statement,
                context=context,
                lake_workdir=lake_workdir,
            )
            results.append(r)
        return results

    # ------------------------------------------------------------------
    # Receipt emission
    # ------------------------------------------------------------------

    def _write_receipt(self, result: ProofResult) -> Path:
        receipt = result.to_receipt()
        ts = receipt["timestamp"].replace(":", "-").split(".")[0]
        fname = f"proof_receipt_{result.theorem[:40]}_{ts}.json"
        path = self.receipt_dir / fname
        path.write_text(json.dumps(receipt, indent=2))
        return path