import os
import json
import subprocess
import sys
import tempfile
import ast
import re
import time
import importlib
import traceback
import signal
from pathlib import Path
from typing import List, Set
import gradio as gr
import huggingface_hub

# Configuration
BOTS_DIR = Path("/app/bots")
DATA_DIR = Path("/app/data")
TEMPLATES_DIR = Path("/app/templates")
CACHE_DIR = Path("/app/cache")

# Create directories
for dir_path in [BOTS_DIR, DATA_DIR, TEMPLATES_DIR, CACHE_DIR]:
    dir_path.mkdir(exist_ok=True)

# Security settings
ALLOWED_BUILTINS = {
    'abs', 'all', 'any', 'ascii', 'bin', 'bool', 'bytearray', 'bytes', 
    'chr', 'complex', 'dict', 'divmod', 'enumerate', 'filter', 'float', 
    'format', 'frozenset', 'hex', 'int', 'iter', 'len', 'list', 'map', 
    'max', 'min', 'next', 'oct', 'ord', 'pow', 'range', 'repr', 'reversed', 
    'round', 'set', 'slice', 'sorted', 'str', 'sum', 'tuple', 'zip',
    'print', 'type', 'isinstance', 'issubclass', 'hasattr', 'getattr'
}

FORBIDDEN_IMPORTS = {
    'os', 'subprocess', 'socket', 'urllib', 
    'importlib', '__builtins__', 'eval', 'exec', 'compile',
    'open', 'file', 'input', 'raw_input'
}

# Package installation tracker
installed_packages = set()

def parse_imports(code: str) -> Set[str]:
    """Parse Python code to find all imports"""
    imports = set()
    try:
        tree = ast.parse(code)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.add(alias.name.split('.')[0])
            elif isinstance(node, ast.ImportFrom):
                if node.module:
                    imports.add(node.module.split('.')[0])
    except SyntaxError:
        import_pattern = r'^(?:from\s+(\w+)|import\s+(\w+))'
        for line in code.split('\n'):
            match = re.match(import_pattern, line.strip())
            if match:
                pkg = match.group(1) or match.group(2)
                if pkg:
                    imports.add(pkg.split('.')[0])
    return imports

def install_package(package_name: str):
    """Install a Python package using pip"""
    global installed_packages
    
    if package_name in installed_packages:
        return True, f"✅ Package {package_name} already installed"
    
    try:
        importlib.import_module(package_name)
        installed_packages.add(package_name)
        return True, f"✅ Package {package_name} already available"
    except ImportError:
        pass
    
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", "--quiet", "--user", package_name],
            capture_output=True,
            text=True,
            timeout=60
        )
        
        if result.returncode == 0:
            installed_packages.add(package_name)
            return True, f"✅ Successfully installed {package_name}"
        else:
            return False, f"❌ Failed to install {package_name}"
    except subprocess.TimeoutExpired:
        return False, f"⏰ Timeout installing {package_name}"
    except Exception as e:
        return False, f"❌ Error: {str(e)[:100]}"

def install_required_packages(code: str) -> List[str]:
    """Automatically install all packages imported in the code"""
    imports = parse_imports(code)
    results = []
    
    for imp in imports:
        if imp in FORBIDDEN_IMPORTS:
            results.append(f"⚠️ Blocked: {imp}")
            continue
        
        if imp not in ALLOWED_BUILTINS:
            success, message = install_package(imp)
            results.append(message)
    
    return results

def run_bot_with_timeout(code, input_data, timeout_seconds=30):
    """Run bot code with timeout using subprocess (avoids thread limits)"""
    
    # Create a temporary script file
    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
        # Build the complete script with the bot code and runner
        script_content = f'''
import sys
import json
import traceback

{code}

try:
    # Install packages check
    install_results = {install_required_packages(code)!r}
    
    # Create bot instance
    bot = Bot()
    
    # Run with input
    input_data = {json.dumps(input_data) if not isinstance(input_data, dict) else json.dumps(input_data)}
    if isinstance(input_data, dict):
        result = bot.run(**input_data)
    else:
        result = bot.run(input_data)
    
    # Output result
    if install_results:
        result = f"[Package Installation]\\n" + "\\n".join(install_results) + f"\\n\\n[Bot Output]\\n{{result}}"
    
    print(json.dumps({{"success": True, "result": result}}))
except Exception as e:
    print(json.dumps({{"success": False, "error": str(e), "traceback": traceback.format_exc()}}))
'''
        f.write(script_content)
        script_path = f.name
    
    try:
        # Run the script with timeout using subprocess
        result = subprocess.run(
            [sys.executable, script_path],
            capture_output=True,
            text=True,
            timeout=timeout_seconds
        )
        
        # Parse the output
        try:
            # Find the JSON output (last line should be JSON)
            output_lines = result.stdout.strip().split('\n')
            json_line = None
            for line in reversed(output_lines):
                if line.startswith('{') and line.endswith('}'):
                    json_line = line
                    break
            
            if json_line:
                output = json.loads(json_line)
                if output.get("success"):
                    return output.get("result", "No result")
                else:
                    return f"❌ Error: {output.get('error', 'Unknown')}\n\n{output.get('traceback', '')}"
            else:
                return f"❌ Error: {result.stderr if result.stderr else result.stdout}"
        except json.JSONDecodeError:
            return f"❌ Error parsing output: {result.stdout}"
            
    except subprocess.TimeoutExpired:
        return f"⏰ Bot execution timed out after {timeout_seconds} seconds"
    except Exception as e:
        return f"❌ Error running bot: {str(e)}"
    finally:
        # Clean up temp file
        try:
            os.unlink(script_path)
        except:
            pass

def execute_in_sandbox(code: str, input_data, timeout_seconds=30, memory_limit_mb=512):
    """Execute bot code in a restricted sandbox environment"""
    return run_bot_with_timeout(code, input_data, timeout_seconds)

# Bot templates
BOT_TEMPLATES = {
    "sentiment_analyzer": {
        "name": "Sentiment Analyzer",
        "description": "Analyzes sentiment of text input",
        "code": '''from transformers import pipeline

class Bot:
    def __init__(self):
        self.sentiment_pipeline = pipeline("sentiment-analysis")
    
    def run(self, text):
        result = self.sentiment_pipeline(text)[0]
        return f"Sentiment: {result['label']} (confidence: {result['score']:.2f})"
'''
    },
    "text_generator": {
        "name": "Text Generator",
        "description": "Generates text based on prompts",
        "code": '''from transformers import pipeline

class Bot:
    def __init__(self):
        self.generator = pipeline("text-generation", model="gpt2")
    
    def run(self, prompt, max_length=100):
        result = self.generator(prompt, max_length=max_length, num_return_sequences=1)
        return result[0]['generated_text']
'''
    },
    "question_answering": {
        "name": "Question Answering",
        "description": "Answers questions based on context",
        "code": '''from transformers import pipeline

class Bot:
    def __init__(self):
        self.qa_pipeline = pipeline("question-answering")
    
    def run(self, question, context):
        result = self.qa_pipeline(question=question, context=context)
        return f"Answer: {result['answer']}\\nConfidence: {result['score']:.2f}"
'''
    },
    "custom_bot": {
        "name": "Custom Bot Template",
        "description": "Create your own bot!",
        "code": '''import transformers
import torch

class Bot:
    def __init__(self):
        # Initialize your model here
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"Running on: {self.device}")
    
    def run(self, input_text):
        # Implement your bot logic here
        return f"Processed: {input_text}"
'''
    },
    "requests_example": {
        "name": "API Bot Example",
        "description": "Fetches data from an API",
        "code": '''# Note: This requires the 'requests' package which will be auto-installed
import requests

class Bot:
    def __init__(self):
        pass
    
    def run(self, url):
        try:
            response = requests.get(url, timeout=10)
            return f"Status: {response.status_code}\\nContent length: {len(response.text)}"
        except Exception as e:
            return f"Error: {e}"
'''
    }
}

def save_bot(name, code, description=""):
    """Save a bot to the bots directory"""
    bot_dir = BOTS_DIR / name
    bot_dir.mkdir(exist_ok=True)
    
    install_results = install_required_packages(code)
    
    with open(bot_dir / "bot.py", "w") as f:
        f.write(code)
    
    metadata = {
        "name": name,
        "description": description,
        "created": time.time(),
        "packages": list(parse_imports(code))
    }
    with open(bot_dir / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)
    
    return f"✅ Bot '{name}' saved!\n\n" + "\n".join(install_results)

def list_bots():
    bots = []
    for bot_dir in BOTS_DIR.iterdir():
        if bot_dir.is_dir():
            meta = bot_dir / "metadata.json"
            if meta.exists():
                with open(meta) as f:
                    bots.append(json.load(f))
    return bots

def load_bot(name):
    bot_dir = BOTS_DIR / name
    if not bot_dir.exists():
        return None, None
    with open(bot_dir / "bot.py") as f:
        code = f.read()
    meta = bot_dir / "metadata.json"
    metadata = json.load(open(meta)) if meta.exists() else {"name": name}
    return code, metadata

def delete_bot(name):
    import shutil
    bot_dir = BOTS_DIR / name
    if bot_dir.exists():
        shutil.rmtree(bot_dir)
        return f"✅ Deleted {name}"
    return f"❌ Not found"

def create_ui():
    with gr.Blocks(title="Python AI Bot Studio", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🤖 Python AI Bot Studio")
        
        with gr.Tabs():
            with gr.TabItem("Create Bot"):
                with gr.Row():
                    with gr.Column():
                        bot_name = gr.Textbox(label="Bot Name")
                        bot_description = gr.Textbox(label="Description")
                        template = gr.Dropdown(
                            choices=[(v["name"], k) for k, v in BOT_TEMPLATES.items()],
                            label="Template",
                            value="custom_bot"
                        )
                        bot_code = gr.Code(
                            label="Bot Code",
                            language="python",
                            value=BOT_TEMPLATES["custom_bot"]["code"],
                            lines=15
                        )
                        check_btn = gr.Button("📦 Check Packages")
                        check_output = gr.Textbox(label="Packages", lines=3)
                        save_btn = gr.Button("💾 Save Bot")
                        save_status = gr.Textbox(label="Status")
                    
                    with gr.Column():
                        test_input = gr.Textbox(label="Test Input", lines=3)
                        test_btn = gr.Button("🧪 Test Bot")
                        test_output = gr.Textbox(label="Output", lines=10)
            
            with gr.TabItem("My Bots"):
                bots_list = gr.Dataframe(headers=["Name", "Description", "Packages", "Created"])
                refresh_btn = gr.Button("🔄 Refresh")
                load_name = gr.Textbox(label="Bot Name")
                load_btn = gr.Button("📂 Load")
                delete_btn = gr.Button("🗑️ Delete")
                loaded_code = gr.Code(language="python", lines=10)
        
        def update_template(tid):
            return BOT_TEMPLATES.get(tid, BOT_TEMPLATES["custom_bot"])["code"]
        
        def check_packages(code):
            imports = parse_imports(code)
            results = []
            for imp in imports:
                if imp in FORBIDDEN_IMPORTS:
                    results.append(f"⚠️ {imp} - BLOCKED")
                elif imp in ALLOWED_BUILTINS:
                    results.append(f"✅ {imp} - Built-in")
                else:
                    results.append(f"📦 {imp} - Will be auto-installed")
            return "\n".join(results) or "No external packages"
        
        template.change(update_template, [template], [bot_code])
        check_btn.click(check_packages, [bot_code], [check_output])
        save_btn.click(save_bot, [bot_name, bot_code, bot_description], [save_status])
        test_btn.click(execute_in_sandbox, [bot_code, test_input], [test_output])
        
        def refresh():
            bots = list_bots()
            return [[b["name"], b.get("description", "")[:30], ", ".join(b.get("packages", [])[:2]), 
                     time.strftime('%H:%M', time.localtime(b["created"]))] for b in bots]
        
        refresh_btn.click(refresh, None, [bots_list])
        load_btn.click(load_bot, [load_name], [loaded_code, save_status])
        delete_btn.click(delete_bot, [load_name], [save_status])
        
        demo.load(refresh, None, [bots_list])
    
    return demo

if __name__ == "__main__":
    demo = create_ui()
    demo.launch(server_name="0.0.0.0", server_port=7860)