"""Small provider-neutral proof repair loop. Dataset statements are trusted input."""
import json
import os
import re
import subprocess
import tempfile
import time
import urllib.request
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Protocol

class Provider(Protocol):
    def complete(self, messages: list[dict[str, str]]) -> str: ...

class HTTPProvider:
    """OpenAI-compatible chat completions, including local serving endpoints."""
    def __init__(self, model, base_url, api_key_env="LLM_API_KEY", timeout=120):
        self.model, self.base_url = model, base_url.rstrip('/')
        self.api_key_env, self.timeout = api_key_env, timeout

    def complete(self, messages):
        headers = {"Content-Type": "application/json"}
        key = os.environ.get(self.api_key_env)
        if key:
            headers["Authorization"] = "Bearer " + key
        request = urllib.request.Request(self.base_url + '/chat/completions',
            data=json.dumps({"model": self.model, "messages": messages}).encode(), headers=headers)
        with urllib.request.urlopen(request, timeout=self.timeout) as response:
            content = json.load(response)["choices"][0]["message"]["content"]
        if not isinstance(content, str):
            raise ValueError("Provider returned non-text content")
        return content

class CommandProvider:
    """Any model: executable receives {messages: [...]} on stdin, returns proof text."""
    def __init__(self, argv, timeout=120):
        if not argv: raise ValueError("Empty provider command")
        self.argv, self.timeout = argv, timeout

    def complete(self, messages):
        proc = subprocess.run(self.argv, input=json.dumps({"messages": messages}),
                              text=True, capture_output=True, timeout=self.timeout)
        if proc.returncode: raise RuntimeError("Provider command failed")
        return proc.stdout

@dataclass
class Check:
    success: bool
    diagnostics: str

class LeanVerifier:
    def __init__(self, command=None, cwd=None, timeout=30):
        self.command = command or ["lean"]
        self.cwd = str(Path(cwd).resolve()) if cwd else None
        self.timeout = timeout

    def check(self, statement, proof, prelude=""):
        # Only a tactic body is accepted. Fail closed on holes and command injection.
        if re.search(r'\b(sorry|admit|axiom|unsafe|implemented_by|run_tac|elab|macro|syntax)\b|--|/\*|/\-|\-/|#', proof):
            return Check(False, "Forbidden proof token (holes, metaprogramming, comments or commands)")
        if not proof.strip(): return Check(False, "Empty proof")
        source = prelude + '\nset_option maxHeartbeats 200000\n'
        source += 'theorem agent_target : ' + statement + ' := by\n'
        source += '\n'.join('  ' + line for line in proof.splitlines())
        source += '\n#print axioms agent_target\n'
        with tempfile.TemporaryDirectory(prefix="lean-check-", dir=self.cwd) as directory:
            path = Path(directory) / 'Candidate.lean'
            path.write_text(source)
            try:
                result = subprocess.run([*self.command, str(path)], cwd=self.cwd,
                    capture_output=True, text=True, timeout=self.timeout)
            except subprocess.TimeoutExpired:
                return Check(False, "Lean timeout")
            diagnostics = result.stdout + result.stderr
        # Reject warnings and require the axiom report; allow only Lean's standard axioms.
        reports = re.findall(r"depends on axioms:\s*\[([^\]]*)\]", diagnostics)
        clean = "does not depend on any axioms" in diagnostics
        if reports:
            clean = all(set(x.strip() for x in report.split(',') if x.strip()) <=
                        {"propext", "Classical.choice", "Quot.sound"} for report in reports)
        return Check(result.returncode == 0 and clean and 'warning:' not in diagnostics,
                     diagnostics[-16000:])

def extract(text):
    match = re.search(r'```(?:lean4?|text)?\s*\n(.*?)```', text, re.S)
    proof = (match.group(1) if match else text).strip()
    if proof == 'by': return ''
    if proof.startswith('by\n') or proof.startswith('by '): proof = proof[2:].strip()
    return proof

def solve(provider, verifier, task, max_attempts=3):
    if max_attempts < 1: raise ValueError("max_attempts must be positive")
    start = time.monotonic()
    messages = [{"role": "system", "content": "Prove the Lean 4 theorem. Return ONLY the tactic body after by. No theorem declaration, comments, sorry, admit, metaprogramming or extra commands."},
                {"role": "user", "content": task.get('prelude', '') + '\n⊢ ' + task['statement']}]
    attempts = []
    for _ in range(max_attempts):
        try:
            proof = extract(provider.complete(messages))
            check = verifier.check(task['statement'], proof, task.get('prelude', ''))
        except Exception as exc:
            # Do not serialize exception bodies: providers may include credentials.
            return {"id": task['id'], "success": False, "error": type(exc).__name__,
                    "attempts": attempts, "seconds": time.monotonic() - start}
        attempts.append({"proof": proof, **asdict(check)})
        if check.success: break
        messages.extend([{"role": "assistant", "content": proof},
                         {"role": "user", "content": "Lean rejected this proof. Repair it:\n" + check.diagnostics}])
    return {"id": task['id'], "success": attempts[-1]['success'],
            "attempts": attempts, "seconds": time.monotonic() - start}