"""
VeloxAI - Fast and Simple AI Library
Text Generation (Vtext) and Image Generation (Vimg)
"""

import requests
import os
import base64
import mimetypes
import json
from typing import Optional, Dict, Any, Iterator, List


class Vtext:
    """Text generation and chat functionality."""
    
    def __init__(self, token: str, base_url: str = "https://bot.insta-acc-sec.workers.dev"):
        """
        Initialize Vtext client.
        
        Args:
            token: API authentication token
            base_url: Base URL for the API
        """
        self.base_url = base_url
        self.headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }
        self._timeout = 300

    def models(self) -> Dict[str, Any]:
        """
        Get available text models.
        
        Returns:
            Dictionary with available models
        """
        try:
            response = requests.get(
                f"{self.base_url}/api/models",
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def chat(
        self,
        message: str,
        model: str = "perplexity-ai",
        system: Optional[str] = None,
        file: Optional[str] = None,
        stream: bool = False
    ) -> Any:
        """
        Send a chat message.
        
        Args:
            message: The message text
            model: Model ID to use
            system: Optional system prompt
            file: Optional file path to attach
            stream: Enable streaming response
            
        Returns:
            Response dict or iterator if streaming
        """
        data = {
            "message": message,
            "model": model,
            "systemPrompt": system
        }

        if file:
            file_data = self._encode_file(file)
            if "error" in file_data:
                return file_data
            data["fileData"] = file_data["encoded"]

        if stream:
            return self._chat_stream(data)
        else:
            return self._chat_complete(data)

    def _chat_complete(self, data: dict) -> Dict[str, Any]:
        """Non-streaming chat request."""
        try:
            response = requests.post(
                f"{self.base_url}/api/v1/chat",
                json=data,
                headers=self.headers,
                timeout=self._timeout
            )
            
            if 'application/json' in response.headers.get('Content-Type', ''):
                return response.json()
            else:
                return {
                    "success": False,
                    "error": f"Server Error: {response.status_code}"
                }
        except requests.exceptions.Timeout:
            return {"success": False, "error": "Request timed out"}
        except Exception as e:
            return {"success": False, "error": str(e)}

    def _chat_stream(self, data: dict) -> Iterator[str]:
        """Streaming chat request."""
        try:
            response = requests.post(
                f"{self.base_url}/api/v1/chat",
                json=data,
                headers=self.headers,
                stream=True,
                timeout=self._timeout
            )
            
            for line in response.iter_lines():
                if line:
                    decoded = line.decode('utf-8')
                    if decoded.startswith('data: '):
                        chunk = decoded[6:]
                        if chunk == '[DONE]':
                            break
                        try:
                            json_chunk = json.loads(chunk)
                            if 'content' in json_chunk:
                                yield json_chunk['content']
                            elif 'delta' in json_chunk:
                                yield json_chunk['delta']
                        except json.JSONDecodeError:
                            yield chunk
        except Exception as e:
            yield f"Error: {str(e)}"

    def completion(
        self,
        prompt: str,
        model: str = "perplexity-ai",
        max_tokens: int = 1000,
        temperature: float = 0.7,
        stream: bool = False
    ) -> Any:
        """
        OpenAI-style completion.
        
        Args:
            prompt: Text prompt
            model: Model to use
            max_tokens: Maximum tokens to generate
            temperature: Sampling temperature (0-1)
            stream: Enable streaming
            
        Returns:
            Completion response or iterator
        """
        return self.chat(
            message=prompt,
            model=model,
            stream=stream
        )

    def train(
        self,
        identity: str,
        role: str,
        extra: str = ""
    ) -> Dict[str, Any]:
        """
        Set fine-tuning/training.
        
        Args:
            identity: AI identity
            role: AI role
            extra: Extra instructions
            
        Returns:
            Success response
        """
        data = {"identity": identity, "role": role, "extra": extra}
        try:
            response = requests.post(
                f"{self.base_url}/user/training",
                json=data,
                headers=self.headers,
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def get_training(self) -> Dict[str, Any]:
        """Get current training settings."""
        try:
            response = requests.get(
                f"{self.base_url}/user/training",
                headers=self.headers,
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def clear_training(self) -> Dict[str, Any]:
        """Clear training settings."""
        try:
            response = requests.delete(
                f"{self.base_url}/user/training",
                headers=self.headers,
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def _encode_file(self, file_path: str) -> Dict[str, Any]:
        """Encode file to base64."""
        if not os.path.exists(file_path):
            return {"error": f"File not found: {file_path}"}

        mime_type, _ = mimetypes.guess_type(file_path)
        if mime_type is None:
            mime_type = 'application/octet-stream'

        try:
            with open(file_path, "rb") as f:
                encoded = base64.b64encode(f.read()).decode('utf-8')
                return {"encoded": f"data:{mime_type};base64,{encoded}"}
        except Exception as e:
            return {"error": f"File error: {e}"}


class Vimg:
    """Image generation functionality."""
    
    def __init__(self, token: str, base_url: str = "https://bot.insta-acc-sec.workers.dev"):
        """
        Initialize Vimg client.
        
        Args:
            token: API authentication token
            base_url: Base URL for the API
        """
        self.base_url = base_url
        self.headers = {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        }

    def generate(
        self,
        prompt: str,
        init_image: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Generate an image.
        
        Args:
            prompt: Text description
            init_image: Optional path to init image (img2img)
            
        Returns:
            Response with record_id
        """
        data = {"prompt": prompt}

        if init_image:
            file_data = self._encode_file(init_image)
            if "error" in file_data:
                return file_data
            data["image"] = file_data["encoded"]

        try:
            response = requests.post(
                f"{self.base_url}/api/v1/image/generate",
                json=data,
                headers=self.headers,
                timeout=60
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def status(self, record_id: str) -> Dict[str, Any]:
        """
        Check generation status.
        
        Args:
            record_id: ID from generate()
            
        Returns:
            Status and images when complete
        """
        try:
            response = requests.get(
                f"{self.base_url}/api/v1/image/status/{record_id}",
                headers=self.headers,
                timeout=30
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"success": False, "error": str(e)}

    def wait(
        self,
        record_id: str,
        timeout: int = 180,
        interval: int = 3
    ) -> Dict[str, Any]:
        """
        Wait for image generation to complete.
        
        Args:
            record_id: ID from generate()
            timeout: Max wait time in seconds
            interval: Polling interval in seconds
            
        Returns:
            Final status with images or error
        """
        import time
        
        elapsed = 0
        while elapsed < timeout:
            status = self.status(record_id)
            
            if status.get('status') == 'DONE':
                return status
            elif status.get('status') == 'FAILED':
                return status
            
            time.sleep(interval)
            elapsed += interval
        
        return {
            "success": False,
            "error": "Timeout waiting for generation"
        }

    def create(
        self,
        prompt: str,
        init_image: Optional[str] = None,
        wait: bool = True
    ) -> Dict[str, Any]:
        """
        Generate and optionally wait for completion.
        
        Args:
            prompt: Text description
            init_image: Optional init image
            wait: Wait for completion
            
        Returns:
            Images if wait=True, record_id if wait=False
        """
        result = self.generate(prompt, init_image)
        
        if not result.get('success'):
            return result
        
        if wait:
            return self.wait(result['record_id'])
        else:
            return result

    def _encode_file(self, file_path: str) -> Dict[str, Any]:
        """Encode image file."""
        if not os.path.exists(file_path):
            return {"error": f"File not found: {file_path}"}

        mime_type, _ = mimetypes.guess_type(file_path)
        if mime_type is None:
            mime_type = 'image/jpeg'

        try:
            with open(file_path, "rb") as f:
                encoded = base64.b64encode(f.read()).decode('utf-8')
                return {"encoded": f"data:{mime_type};base64,{encoded}"}
        except Exception as e:
            return {"error": f"File error: {e}"}


__version__ = "1.0.0"
__all__ = ["Vtext", "Vimg"]