"""
iDEF Geometry API - Python SDK
Official Python client for the iDEF Geometry theorem proving API.

Installation:
    pip install requests

Usage:
    from idef_geometry import GeometryAPI
    
    api = GeometryAPI("gai_live_your_key")
    result = api.prove("Triangle ABC with AB=AC. Prove ABC=ACB.")
    print(result.proved)
"""

import requests
import time
from dataclasses import dataclass
from typing import Optional, Literal

__version__ = "1.0.0"

API_BASE = "https://idef-mathematics.com/api/v1"


@dataclass
class ProofResult:
    """Result of a geometry proof attempt"""
    proved: Optional[bool]
    proof_trace: str
    computation_time: float
    model_used: str
    job_id: str
    
    def __repr__(self):
        status = "✓ PROVED" if self.proved else ("✗ NOT PROVED" if self.proved is False else "? UNKNOWN")
        return f"ProofResult({status}, time={self.computation_time}s, model={self.model_used})"


class GeometryAPIError(Exception):
    """API Error"""
    def __init__(self, message: str, code: str = None, status_code: int = None):
        super().__init__(message)
        self.code = code
        self.status_code = status_code


class GeometryAPI:
    """
    iDEF Geometry API Client
    
    Args:
        api_key: Your API key (get one at https://idef-mathematics.com/api)
        base_url: API base URL (default: https://idef-mathematics.com/api/v1)
        timeout: Request timeout in seconds (default: 30)
    
    Example:
        api = GeometryAPI("gai_live_xxxxx")
        result = api.prove("Triangle ABC with AB=AC. Prove ABC=ACB.", model="A5")
        print(result.proved)
        print(result.proof_trace)
    """
    
    def __init__(self, api_key: str, base_url: str = API_BASE, timeout: int = 30):
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = requests.Session()
        self.session.headers.update({
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "User-Agent": f"idef-geometry-python/{__version__}"
        })
    
    def prove(
        self,
        problem: str,
        model: Literal["A5", "A6"] = "A5",
        timeout: int = 120,
        poll_interval: float = 1.0,
        max_wait: int = 300
    ) -> ProofResult:
        """
        Prove a geometry theorem.
        
        Args:
            problem: The geometry problem to prove (natural language or GIVEN/GOAL format)
            model: Model to use - "A5" (Wu's Method) or "A6" (Gröbner Basis)
            timeout: Max computation time in seconds
            poll_interval: How often to poll for results
            max_wait: Maximum time to wait for result
        
        Returns:
            ProofResult with proved status and proof trace
        
        Raises:
            GeometryAPIError: If API call fails
        """
        # Submit problem
        response = self.session.post(
            f"{self.base_url}/problems",
            json={
                "problem": problem,
                "options": {"model": model, "timeout": timeout}
            },
            timeout=self.timeout
        )
        
        if not response.ok:
            error = response.json()
            raise GeometryAPIError(
                error.get("message", response.text),
                error.get("code"),
                response.status_code
            )
        
        data = response.json()
        job_id = data["job_id"]
        
        # Poll for result
        start_time = time.time()
        while time.time() - start_time < max_wait:
            result_response = self.session.get(
                f"{self.base_url}/jobs/{job_id}",
                timeout=self.timeout
            )
            
            if not result_response.ok:
                raise GeometryAPIError(f"Failed to get job status: {result_response.text}")
            
            result = result_response.json()
            status = result.get("status")
            
            if status == "completed":
                r = result.get("result", {})
                return ProofResult(
                    proved=r.get("proved"),
                    proof_trace=r.get("proof_trace", ""),
                    computation_time=r.get("computation_time", 0),
                    model_used=r.get("model_used", model),
                    job_id=job_id
                )
            elif status == "failed":
                raise GeometryAPIError(result.get("error", "Job failed"))
            elif status == "timeout":
                raise GeometryAPIError("Computation timed out")
            
            time.sleep(poll_interval)
        
        raise GeometryAPIError(f"Timed out waiting for result after {max_wait}s")
    
    def submit(
        self,
        problem: str,
        model: Literal["A5", "A6"] = "A5",
        timeout: int = 120
    ) -> str:
        """
        Submit a problem without waiting for result. Returns job_id.
        Use get_result() to poll for completion.
        """
        response = self.session.post(
            f"{self.base_url}/problems",
            json={
                "problem": problem,
                "options": {"model": model, "timeout": timeout}
            },
            timeout=self.timeout
        )
        
        if not response.ok:
            error = response.json()
            raise GeometryAPIError(error.get("message", response.text))
        
        return response.json()["job_id"]
    
    def get_result(self, job_id: str) -> dict:
        """Get the status/result of a submitted job"""
        response = self.session.get(
            f"{self.base_url}/jobs/{job_id}",
            timeout=self.timeout
        )
        
        if not response.ok:
            raise GeometryAPIError(f"Failed to get job: {response.text}")
        
        return response.json()
    
    def models(self) -> list:
        """List available models"""
        response = self.session.get(f"{self.base_url}/models", timeout=self.timeout)
        if not response.ok:
            raise GeometryAPIError(f"Failed to list models: {response.text}")
        return response.json()
    
    def status(self) -> dict:
        """Get API status"""
        response = self.session.get(f"{self.base_url}/status", timeout=self.timeout)
        if not response.ok:
            raise GeometryAPIError(f"Failed to get status: {response.text}")
        return response.json()


# Convenience function
def prove(api_key: str, problem: str, model: str = "A5") -> ProofResult:
    """Quick prove without creating client instance"""
    return GeometryAPI(api_key).prove(problem, model=model)


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 3:
        print("Usage: python idef_geometry.py <API_KEY> <PROBLEM> [MODEL]")
        print("Example: python idef_geometry.py gai_free_xxx 'Triangle ABC...' A5")
        sys.exit(1)
    
    api_key = sys.argv[1]
    problem = sys.argv[2]
    model = sys.argv[3] if len(sys.argv) > 3 else "A5"
    
    print(f"Proving with model {model}...")
    result = prove(api_key, problem, model)
    
    print(f"\n{'='*60}")
    print(result)
    print(f"{'='*60}")
    print(result.proof_trace)