import requests
import time
import json

class APIClient:
    def __init__(self, base_url="http://127.0.0.1:8000"):
        self.base_url = base_url
        self.endpoint = "/generate_plan"

    def send_debug_stage(
        self, prompt: str, drone_state: str = "on_ground", target_stage: int = 1, timeout=120
    ):
        """
        Stage 分阶段调试请求。
        target_stage: 1-6
        Returns: 同 send_request 结构，data 为 debug 端返回的 {target_stage, stage_name, output}
        """
        url = f"{self.base_url}/debug_stage"
        payload = {
            "user_prompt": prompt,
            "drone_state": drone_state,
            "target_stage": target_stage,
        }
        headers = {"Content-Type": "application/json"}
        start_time = time.time()
        try:
            response = requests.post(url, json=payload, headers=headers, timeout=timeout)
            latency = time.time() - start_time
            response.raise_for_status()
            try:
                data = response.json()
                if "error" in data and len(data) == 1:
                    return {
                        "success": False,
                        "data": data,
                        "latency": latency,
                        "error": data["error"],
                        "http_status": response.status_code,
                    }
                return {
                    "success": True,
                    "data": data,
                    "latency": latency,
                    "error": None,
                    "http_status": response.status_code,
                }
            except json.JSONDecodeError:
                return {
                    "success": False,
                    "data": None,
                    "latency": latency,
                    "error": f"Invalid JSON: {response.text[:200]}",
                    "http_status": response.status_code,
                }
        except requests.exceptions.RequestException as e:
            latency = time.time() - start_time
            return {
                "success": False,
                "data": None,
                "latency": latency,
                "error": str(e),
                "http_status": getattr(e.response, "status_code", None) if hasattr(e, "response") and e.response is not None else None,
            }

    def send_request(self, prompt, drone_state="on_ground", timeout=60):
        """
        Sends a request to the API and returns a structured result.
        Returns:
            dict: {
                "success": bool,
                "data": dict or None,
                "latency": float (seconds),
                "error": str or None,
                "http_status": int or None
            }
        """
        url = f"{self.base_url}{self.endpoint}"
        payload = {"user_prompt": prompt, "drone_state": drone_state}
        headers = {"Content-Type": "application/json"}
        
        start_time = time.time()
        try:
            response = requests.post(url, json=payload, headers=headers, timeout=timeout)
            latency = time.time() - start_time
            
            response.raise_for_status()
            
            try:
                data = response.json()
                return {
                    "success": True,
                    "data": data,
                    "latency": latency,
                    "error": None,
                    "http_status": response.status_code
                }
            except json.JSONDecodeError:
                return {
                    "success": False,
                    "data": None,
                    "latency": latency,
                    "error": f"Invalid JSON response: {response.text[:200]}",
                    "http_status": response.status_code
                }
                
        except requests.exceptions.RequestException as e:
            latency = time.time() - start_time
            return {
                "success": False,
                "data": None,
                "latency": latency,
                "error": str(e),
                "http_status": getattr(e.response, 'status_code', None)
            }