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

class CerebrasCodeCompletion:
    def __init__(self):
        self.client = Cerebras(
            api_key="csk-8y84rpd848e2m4rycdk33eyn28ykvrvfcxfj5mdpv8m8wxrh"
        )
    
    async def complete_code(self, context: str, language: str = "python", file_content: str = ""):
        """Generate code + insertion point"""
        try:
            system_prompt = f"""You are an expert {language} developer. 
You will be given a coding request and the full contents of a target file. 
You must return:
- the code block you would generate
- the best place in the file to insert it (by identifying a line of existing code to insert *after*).

Always return a valid JSON object with exactly two keys:
  "insert_after": (a line from the file to insert after),
  "code": (the code block to insert).

ONLY return the JSON, no explanation, no markdown."""
            
            user_prompt = f"""Request: {context}

Target File:
\"\"\"
{file_content}
\"\"\""""

            stream = self.client.chat.completions.create(
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt}
                ],
                model="gpt-oss-120b",
                stream=True,
                max_completion_tokens=1024,
                temperature=0.3,
                top_p=0.9,
                reasoning_effort="medium"
            )

            completion = ""
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    completion += chunk.choices[0].delta.content

            return completion.strip()

        except Exception as e:
            return json.dumps({"error": str(e)})

    async def chat_completion(self, prompt: str):
        try:
            system_prompt = """You are an expert programming assistant..."""
            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: {str(e)}"

# --- CLI Entry ---
async def main():
    if len(sys.argv) < 3:
        print("Usage: python extension.py <command> <args...>")
        return
    
    command = sys.argv[1]
    completion_provider = CerebrasCodeCompletion()

    if command == "complete":
        if len(sys.argv) < 4:
            print("Usage: python extension.py complete <context> <language> [file_content]")
            return

        context = sys.argv[2]
        language = sys.argv[3] if len(sys.argv) > 3 else "python"
        file_content = sys.argv[4] if len(sys.argv) > 4 else ""

        result = await completion_provider.complete_code(context, language, file_content)
        print(result)

    elif command == "chat":
        if len(sys.argv) < 3:
            print("Usage: python extension.py chat <prompt>")
            return

        prompt = sys.argv[2]
        result = await completion_provider.chat_completion(prompt)
        print(result)

    else:
        print(f"Unknown command: {command}")

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