from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class CodeCompletionEngine:
    def __init__(self, model_name: str = "bigcode/starcoder"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=torch.float16, device_map="auto"
        )

    def complete(self, code_context: str, max_new_tokens: int = 100,
                 temperature: float = 0.2) -> str:
        inputs = self.tokenizer(code_context, return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs, max_new_tokens=max_new_tokens,
                temperature=temperature, top_p=0.95
            )
        generated = self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:])
        return generated.split("\n")[0]

    def fill_in_middle(self, prefix: str, suffix: str) -> str:
        prompt = f"<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>"
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=200)
        return self.tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:])

# Usage
engine = CodeCompletionEngine()
completion = engine.complete("def fibonacci(n):\n    if n <= 1:\n        return n\n    return ")

from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate

class CodeAssistant:
    def __init__(self, llm=None):
        self.llm = llm or ChatOpenAI(model="gpt-4", temperature=0)
        self.chat_history = []

    def explain_code(self, code: str, language: str = "python") -> str:
        prompt = PromptTemplate.from_template(
            """Explain this {language} code in detail:
            ```{language}
            {code}
            ```
            Provide a clear explanation of what each part does:"""
        )
        chain = prompt | self.llm
        return chain.invoke({"code": code, "language": language}).content

    def refactor_code(self, code: str, instructions: str = "optimize") -> str:
        prompt = PromptTemplate.from_template(
            """Refactor this code to {instructions}:
            ```python
            {code}
            ```
            Return the improved code:"""
        )
        chain = prompt | self.llm
        return chain.invoke({"code": code, "instructions": instructions}).content

    def debug_code(self, code: str, error: str = "") -> str:
        prompt = PromptTemplate.from_template(
            """Find and fix bugs in this code:
            ```python
            {code}
            ```
            Error: {error}
            Provide the fixed code and explain the bugs:"""
        )
        chain = prompt | self.llm
        return chain.invoke({"code": code, "error": error}).content

    def generate_docs(self, code: str) -> str:
        prompt = PromptTemplate.from_template(
            """Generate comprehensive documentation for this code:
            ```python
            {code}
            ```
            Include docstrings, type hints, and usage examples:"""
        )
        chain = prompt | self.llm
        return chain.invoke({"code": code}).content

    def write_tests(self, code: str, framework: str = "pytest") -> str:
        prompt = PromptTemplate.from_template(
            """Write {framework} tests for this code:
            ```python
            {code}
            ```
            Include edge cases and assertions:"""
        )
        chain = prompt | self.llm
        return chain.invoke({"code": code, "framework": framework}).content

    def chat(self, message: str, code_context: str = "") -> str:
        self.chat_history.append({"role": "user", "content": message})
        prompt = f"""You are a code assistant. Context: {code_context}
        Chat history: {self.chat_history[-5:]}
        User: {message}
        Assistant:"""
        response = self.llm.invoke(prompt).content
        self.chat_history.append({"role": "assistant", "content": response})
        return response

# Usage
assistant = CodeAssistant()
explanation = assistant.explain_code("def quicksort(arr): ...")
refactored = assistant.refactor_code("for i in range(len(arr)): ...", "make it more Pythonic")
tests = assistant.write_tests("def add(a, b): return a + b")

class CodeReviewAgent:
    def __init__(self, llm):
        self.llm = llm
        self.checks = [
            "security", "performance", "readability",
            "best_practices", "error_handling"
        ]

    def review(self, code: str, checks: list = None) -> dict:
        checks = checks or self.checks
        results = {}
        for check in checks:
            results[check] = self._run_check(code, check)
        return results

    def _run_check(self, code: str, check_type: str) -> dict:
        prompts = {
            "security": "Review for security vulnerabilities:",
            "performance": "Review for performance issues:",
            "readability": "Review for code readability:",
            "best_practices": "Review against Python best practices:",
            "error_handling": "Review error handling:"
        }
        prompt = f"""{prompts[check_type]}
        ```python
        {code}
        ```
        Rate 1-10 and list issues:"""

        result = self.llm.invoke(prompt).content
        return {"review": result, "check_type": check}

    def generate_summary(self, reviews: dict) -> str:
        summary = "Code Review Summary:\n"
        for check, result in reviews.items():
            summary += f"- {check}: {result['review'][:100]}...\n"
        return summary

# Usage
reviewer = CodeReviewAgent(llm)
reviews = reviewer.review("def process(data): return eval(data)")
# Flags security issue with eval()