#!/usr/bin/env python3
"""
Detoxify-Medium Model Inference Script

This script provides a simple interface to query the Detoxify-Medium model
for text detoxification. It can handle both streaming and non-streaming responses.

Usage:
    python detoxify_inference.py

Requirements:
    - requests library: pip install requests
    - The Detoxify-Medium model server running on http://127.0.0.1:8000
"""

import requests
import json
import time
import sys
from typing import Optional, List, Dict, Any


class DetoxifyClient:
    """Client for interacting with the Detoxify-Medium model server."""

    def __init__(self, base_url: str = "http://127.0.0.1:8000"):
        self.base_url = base_url
        self.completion_url = f"{base_url}/completion"
        self.health_url = f"{base_url}/health"

    def check_server_health(self) -> bool:
        """Check if the model server is running and healthy."""
        try:
            response = requests.get(self.health_url, timeout=5)
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False

    def detoxify_text(self,
                     text: str,
                     max_tokens: int = 256,
                     temperature: float = 0.7,
                     stream: bool = False) -> str:
        """
        Detoxify the given text using the model.

        Args:
            text: The text to detoxify
            max_tokens: Maximum number of tokens to generate
            temperature: Sampling temperature (0.0 to 1.0)
            stream: Whether to use streaming response

        Returns:
            The detoxified text
        """
        # Format the prompt according to the model's expected format
        prompt = f"Instruction: Rewrite the provided text to remove the toxicity.\n\nInput: {text}\n\nResponse: "

        payload = {
            "prompt": prompt,
            "max_tokens": max_tokens,
            "temperature": temperature,
            "stream": stream
        }

        headers = {
            'Content-Type': 'application/json'
        }

        try:
            if stream:
                return self._handle_streaming_response(payload, headers)
            else:
                return self._handle_non_streaming_response(payload, headers)
        except requests.exceptions.RequestException as e:
            return f"Error: Failed to connect to server - {e}"
        except Exception as e:
            return f"Error: {e}"

    def _handle_non_streaming_response(self, payload: Dict[str, Any], headers: Dict[str, str]) -> str:
        """Handle non-streaming response from the server."""
        response = requests.post(self.completion_url, json=payload, headers=headers, timeout=30)

        if response.status_code == 200:
            result = response.json()
            return result.get('content', 'No content received')
        else:
            return f"Error: Server returned status code {response.status_code}"

    def _handle_streaming_response(self, payload: Dict[str, Any], headers: Dict[str, str]) -> str:
        """Handle streaming response from the server."""
        response = requests.post(self.completion_url, json=payload, headers=headers, stream=True, timeout=30)

        if response.status_code != 200:
            return f"Error: Server returned status code {response.status_code}"

        full_response = ""
        print("Streaming response:")
        print("-" * 50)

        try:
            for line in response.iter_lines():
                if line:
                    line_str = line.decode('utf-8')
                    if line_str.startswith('data: '):
                        data_str = line_str[6:]  # Remove 'data: ' prefix
                        if data_str.strip() == '[DONE]':
                            break
                        try:
                            data = json.loads(data_str)
                            content = data.get('content', '')
                            if content:
                                print(content, end='', flush=True)
                                full_response += content
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            return f"Error during streaming: {e}"

        print("\n" + "-" * 50)
        return full_response


def run_example_prompts():
    """Run a series of example prompts to demonstrate the model's capabilities."""

    # Initialize the client
    client = DetoxifyClient()

    # Check if server is running
    print("🔍 Checking server health...")
    if not client.check_server_health():
        print("❌ Error: Server is not running or not accessible!")
        print("Please start the server first by running:")
        print("  cd /Users/niko/Projects/Minibase/detoxify-medium")
        print("  ./Minibase-detoxify-medium.app/Contents/MacOS/run_server")
        return

    print("✅ Server is running and healthy!")
    print()

    # Example prompts for testing detoxification
    example_prompts = [
        "You are such a retarded piece of shit!",
        "Fuck you and your mother you son of a whore!",
        "Lick my balls. Gobble gobble!",
        "Suck my clit you stupid bitch!",
        "I fucked your mother!",
        "Why are you being such a dumbass!",
        "STFU asshole!",
        "Wah wah go cry you little pussy!",
        "Fucking idiot",
        "Filha da puta"
    ]

    print("🚀 Running detoxification examples...")
    print("=" * 60)

    for i, prompt in enumerate(example_prompts, 1):
        print(f"\n📝 Example {i}:")
        print(f"Original: {prompt}")

        # Get detoxified version
        detoxified = client.detoxify_text(prompt, max_tokens=150, temperature=0.7)

        print(f"Detoxified: {detoxified}")
        print("-" * 40)

        # Small delay to avoid overwhelming the server
        time.sleep(0.5)

    print("\n🎉 All examples completed!")


def interactive_mode():
    """Run in interactive mode where user can input their own text."""
    client = DetoxifyClient()

    # Check if server is running
    print("🔍 Checking server health...")
    if not client.check_server_health():
        print("❌ Error: Server is not running or not accessible!")
        print("Please start the server first by running:")
        print("  cd /Users/niko/Projects/Minibase/detoxify-medium")
        print("  ./Minibase-detoxify-medium.app/Contents/MacOS/run_server")
        return

    print("✅ Server is running and healthy!")
    print("\n🎯 Interactive Detoxification Mode")
    print("Enter text to detoxify (or 'quit' to exit):")
    print("=" * 50)

    while True:
        try:
            user_input = input("\n📝 Enter text: ").strip()

            if user_input.lower() in ['quit', 'exit', 'q']:
                print("👋 Goodbye!")
                break

            if not user_input:
                print("Please enter some text to detoxify.")
                continue

            print("\n🔄 Processing...")

            # Ask user for preferences
            try:
                max_tokens = int(input("Max tokens (default 256): ") or "256")
                temperature = float(input("Temperature 0.0-1.0 (default 0.7): ") or "0.7")
                stream = input("Stream response? (y/N): ").lower().startswith('y')
            except ValueError:
                max_tokens = 256
                temperature = 0.7
                stream = False

            # Get detoxified version
            detoxified = client.detoxify_text(
                user_input,
                max_tokens=max_tokens,
                temperature=temperature,
                stream=stream
            )

            if not stream:
                print(f"\n✨ Detoxified: {detoxified}")

        except KeyboardInterrupt:
            print("\n\n👋 Goodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {e}")


def main():
    """Main function to run the script."""
    print("🧹 Detoxify-Medium Model Inference Script")
    print("=" * 50)

    if len(sys.argv) > 1 and sys.argv[1] == "--interactive":
        interactive_mode()
    else:
        print("Choose an option:")
        print("1. Run example prompts")
        print("2. Interactive mode")
        print("3. Exit")

        while True:
            try:
                choice = input("\nEnter your choice (1-3): ").strip()

                if choice == "1":
                    run_example_prompts()
                    break
                elif choice == "2":
                    interactive_mode()
                    break
                elif choice == "3":
                    print("👋 Goodbye!")
                    break
                else:
                    print("Please enter 1, 2, or 3.")
            except KeyboardInterrupt:
                print("\n\n👋 Goodbye!")
                break


if __name__ == "__main__":
    main()