from typing import List, Optional, Union
import requests
from pydantic import ValidationError

from .exceptions import APIError, GoMemoryError
from .models import (
    FullReqStatus,
    InsertMemoryRequest,
    Memory,
    MemoryInsertionResponse,
    MemoryRetrievalRequest,
    Message,
    UserResponse
)

class GoMemoryClient:
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()

    def _handle_response(self, response: requests.Response):
        """Internal helper to catch APIError structs and raise them as Python exceptions."""
        try:
            data = response.json()
        except ValueError:
            data = response.text

        if not response.ok:
            # Check if it matches the APIError struct structure
            if isinstance(data, dict) and "Status" in data and "Message" in data:
                raise APIError(
                    status=data.get("Status", response.status_code),
                    message=data.get("Message", "Unknown error"),
                    error_details=str(data.get("Error", ""))
                )
            response.raise_for_status()
            
        return data

    def health_check(self) -> str:
        """GET /health"""
        url = f"{self.base_url}/health"
        response = self.session.get(url)
        # Endpoint returns literal string via writeJSON
        return self._handle_response(response)

    def create_user(self) -> UserResponse:
        """POST /create/user"""
        url = f"{self.base_url}/create/user"
        response = self.session.post(url)
        data = self._handle_response(response)
        return UserResponse(**data)

    def add_memory(self, user_id: str, messages: List[Message]) -> MemoryInsertionResponse:
        """POST /add_memory"""
        url = f"{self.base_url}/add_memory"
        
        # Pydantic validation on input
        req_data = InsertMemoryRequest(userId=user_id, messages=messages)
        
        response = self.session.post(
            url, 
            json=req_data.model_dump(by_alias=True)
        )
        data = self._handle_response(response)
        return MemoryInsertionResponse(**data)

    def get_memory(
        self, 
        user_id: str, 
        messages: Optional[List[Message]] = None, 
        query: Optional[str] = None, 
        threshold: Optional[float] = None
    ) -> List[Memory]:
        """POST /get_memory"""
        url = f"{self.base_url}/get_memory"
        
        # Pydantic handles the XOR validation for messages/query here
        req_data = MemoryRetrievalRequest(
            userId=user_id, 
            messages=messages, 
            query=query, 
            threshold=threshold
        )
        
        response = self.session.post(
            url, 
            json=req_data.model_dump(by_alias=True, exclude_none=True)
        )
        data = self._handle_response(response)
        
        # Expected to return a list of memories on success
        if isinstance(data, list):
            return [Memory(**m) for m in data]
        raise GoMemoryError(f"Unexpected response format: {data}")

    def get_status(self, req_id: str) -> FullReqStatus:
        """GET /get_status/{id}"""
        url = f"{self.base_url}/get_status/{req_id}"
        response = self.session.get(url)
        data = self._handle_response(response)
        return FullReqStatus(**data)