import os
import json
from cerebras.cloud.sdk import Cerebras

class CerebrasChatIntegration:
    def __init__(self):
        # Use the provided API key directly
        self.client = Cerebras(
            api_key="csk-8y84rpd848e2m4rycdk33eyn28ykvrvfcxfj5mdpv8m8wxrh"
        )
    
    async def process_chat_message(self, message: str):
        """Process chat messages that start with 'using oss'"""
        try:
            # Check if message starts with "using oss"
            if not message.lower().startswith("using oss"):
                return None
            
            # Extract the actual prompt (remove "using oss" prefix)
            prompt = message[9:].strip()  # Remove "using oss " (9 characters)
            
            if not prompt:
                return "Please provide a prompt after 'using oss'"
            
            # Generate response using Cerebras
            system_prompt = """You are an expert programming assistant. 
            Provide helpful, accurate, and concise responses to programming questions.
            When asked to write code, provide clean, well-documented code with explanations.
            Be conversational and helpful."""
            
            stream = self.client.chat.completions.create(
                messages=[
                    {
                        "role": "system",
                        "content": system_prompt
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                model="gpt-oss-120b",
                stream=True,
                max_completion_tokens=2048,
                temperature=0.7,
                top_p=0.9,
                reasoning_effort="medium"
            )
            
            response = ""
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    response += chunk.choices[0].delta.content
            
            return response.strip()
            
        except Exception as e:
            return f"Error processing request: {str(e)}"

# Global instance for easy access
chat_integration = CerebrasChatIntegration()

# Function that Cursor can call directly
async def handle_chat_message(message: str):
    """Handle chat messages from Cursor - this is the main entry point"""
    return await chat_integration.process_chat_message(message)

# Command line interface
async def main():
    if len(sys.argv) < 2:
        print("Usage: python chat_integration.py <message>")
        return
    
    message = sys.argv[1]
    result = await chat_integration.process_chat_message(message)
    
    if result:
        print(result)
    else:
        print("Message does not start with 'using oss'")

if __name__ == "__main__":
    import sys
    import asyncio
    asyncio.run(main())