import os
import subprocess
import sys
import socket
import datetime
import json
import re

# =====================================================================
# 1. STRICT /px DIRECTORY ARCHITECTURE & AUTO-BOOTSTRAPPER
# =====================================================================
BASE_DIR = "/px"
try:
    os.makedirs(BASE_DIR, exist_ok=True)
except PermissionError:
    BASE_DIR = "./px"
    os.makedirs(BASE_DIR, exist_ok=True)

GLOBAL_CONTEXT_DIR = os.path.join(BASE_DIR, "global_context")
IPS_DIR = os.path.join(BASE_DIR, "ips")
LOG_FILE = os.path.join(BASE_DIR, "server.log")

os.makedirs(GLOBAL_CONTEXT_DIR, exist_ok=True)
os.makedirs(IPS_DIR, exist_ok=True)

global_context_path = os.path.join(GLOBAL_CONTEXT_DIR, "shared_knowledge.md")
if not os.path.exists(global_context_path):
    with open(global_context_path, "w") as f:
        f.write("# Global Context & Workspace\nShared environment memory stored under /px.")

# Auto-install necessary dependencies
required_packages = ["gradio==4.44.1", "requests", "spaces", "torch", "transformers", "accelerate", "beautifulsoup4"]
for package in required_packages:
    try:
        pkg_name = package.split("==")[0]
        __import__(pkg_name)
    except ImportError:
        print(f"📦 Installing missing dependency: {package}...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])

import requests
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import gradio as gr
from bs4 import BeautifulSoup

# =====================================================================
# 2. NETWORK IP LOGGING UNDER /px
# =====================================================================
def get_ip_addresses():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
    except Exception:
        local_ip = "127.0.0.1"
    try:
        public_ip = requests.get("https://api.ipify.org", timeout=5).text.strip()
    except Exception:
        public_ip = "Online Server Space"
    return local_ip, public_ip

local_ip, server_public_ip = get_ip_addresses()
timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
with open(LOG_FILE, "a") as log_file:
    log_file.write(f"[{timestamp}] - Web Chat Online | Server IP: {server_public_ip}\n")

print("\n" + "═" * 65)
print(f"🌐 PX TECH SOLUTIONS - EVAL BETA 0.2 (AGENTIC ZEROGPU CORE)")
print(f"   • Base Directory : {BASE_DIR}")
print(f"   • Server IP      : {server_public_ip}")
print("═" * 65 + "\n")

# =====================================================================
# 3. ZEROGPU-COMPATIBLE MODEL INITIALIZATION
# =====================================================================
MODEL_ID = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
print(f"🧠 Loading AI Model ({MODEL_ID})...")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# In ZeroGPU, device_map="auto" allows the framework to dynamically shift tensors
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# =====================================================================
# 4. ROBUST AGENTIC TOOLS SUITE
# =====================================================================
def execute_tool(name, args, ip_folder, yolo_mode):
    try:
        if name == "bash":
            cmd = args.get("command", "")
            if not yolo_mode and any(danger in cmd for danger in ["rm -rf /", "mkfs", "dd if="]):
                return "❌ Security Error: Command blocked by YOLO Security Sandbox."
            res = subprocess.run(cmd, shell=True, capture_output=True, text=True, cwd=ip_folder, timeout=45)
            output = res.stdout if res.returncode == 0 else f"STDOUT:\n{res.stdout}\nSTDERR:\n{res.stderr}"
            return output[:8000] if output else "✅ Command executed successfully (No output)."

        elif name == "write_file":
            filepath = args.get("filepath")
            content = args.get("content")
            safe_path = os.path.join(ip_folder, filepath) if not os.path.isabs(filepath) else filepath
            os.makedirs(os.path.dirname(safe_path), exist_ok=True)
            with open(safe_path, "w", encoding="utf-8") as f:
                f.write(content)
            return f"✅ File successfully written to {safe_path}"

        elif name == "read_file":
            filepath = args.get("filepath")
            safe_path = os.path.join(ip_folder, filepath) if not os.path.isabs(filepath) else filepath
            if not os.path.exists(safe_path) and os.path.exists(filepath):
                safe_path = filepath
            if not os.path.exists(safe_path):
                return f"❌ Error: File '{filepath}' not found."
            with open(safe_path, "r", encoding="utf-8", errors="ignore") as f:
                return f.read()[:15000]

        elif name == "web_search":
            query = args.get("query")
            url = f"https://html.duckduckgo.com/html/?q={requests.utils.quote(query)}"
            headers = {"User-Agent": "Mozilla/5.0"}
            resp = requests.get(url, headers=headers, timeout=10)
            soup = BeautifulSoup(resp.text, "html.parser")
            results = [a.get_text(strip=True) for a in soup.find_all("a", class_="result__snippet", limit=5)]
            return "\n\n".join(results) if results else "❌ No web search results retrieved."

        else:
            return f"❌ Error: Unknown tool '{name}'"
    except Exception as e:
        return f"❌ Tool execution failed: {str(e)}"

# =====================================================================
# 5. MULTI-TURN AGENTIC INFERENCE LOOP (ZeroGPU Decorated)
# =====================================================================
@spaces.GPU(duration=120)
def predict(message, history, system_prompt, temperature, yolo_mode, request: gr.Request):
    # Extract & isolate client IP
    client_ip = "127.0.0.1"
    if request and request.headers:
        client_ip = request.headers.get("x-forwarded-for", client_ip).split(",")[0].strip()
    elif request and request.client:
        client_ip = request.client.host

    safe_ip = client_ip.replace(":", "_").replace(".", "_")
    ip_folder = os.path.join(IPS_DIR, safe_ip)
    os.makedirs(ip_folder, exist_ok=True)
    history_log_path = os.path.join(ip_folder, "chat_history.log")

    # Global context & YOLO setup
    security_mode = "DISABLED (GOD-MODE)" if yolo_mode else "ENABLED (SECURE)"
    global_text = open(global_context_path, "r").read() if os.path.exists(global_context_path) else ""

    agent_instructions = """
You have access to powerful tools. To use a tool, you MUST output a JSON block wrapped in <tool_call> tags:

<tool_call>
{"name": "bash", "arguments": {"command": "ls -la"}}
</tool_call>

Available Tools:
1. bash - Execute shell commands. Args: {"command": "<command string>"}
2. write_file - Write to a file. Args: {"filepath": "<path>", "content": "<string>"}
3. read_file - Read a file. Args: {"filepath": "<path>"}
4. web_search - Search the live web. Args: {"query": "<search query>"}

If you use a tool, wait for the <tool_response> before continuing.
"""

    if yolo_mode:
        yolo_instructions = "\n🔥 [YOLO MODE: ON] Absolute root permissions granted. Execute shell scripts and downloads freely under /px."
    else:
        yolo_instructions = "\n🔒 [YOLO MODE: OFF] Standard security limits apply."

    enhanced_system_prompt = f"{system_prompt}\n{agent_instructions}\n{yolo_instructions}\n\n--- [GLOBAL CONTEXT] ---\n{global_text}\n\n--- [USER IP CONTEXT] ---\nClient IP: {client_ip} | Folder: {ip_folder} | Security: {security_mode}"

    # Build chat history
    chat_messages = [{"role": "system", "content": enhanced_system_prompt}]
    for human, assistant in history:
        chat_messages.append({"role": "user", "content": human})
        chat_messages.append({"role": "assistant", "content": assistant})
    chat_messages.append({"role": "user", "content": message})

    with open(history_log_path, "a") as hf:
        hf.write(f"[{timestamp}] User (YOLO: {yolo_mode}): {message}\n")

    # Autonomous Multi-Step Loop
    max_turns = 4
    turn = 0
    accumulated_output = ""

    while turn < max_turns:
        turn += 1
        prompt = tokenizer.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=1024,
            temperature=max(temperature, 0.01),
            do_sample=True if temperature > 0 else False
        )

        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()

        turn_response = ""
        for new_text in streamer:
            turn_response += new_text
            yield accumulated_output + turn_response

        accumulated_output += turn_response
        chat_messages.append({"role": "assistant", "content": turn_response})

        # Intercept and Parse Tool Calls
        tool_call_match = re.search(r'<tool_call>(.*?)</tool_call>', turn_response, re.DOTALL)
        if tool_call_match:
            try:
                tool_json = json.loads(tool_call_match.group(1).strip())
                func_name = tool_json.get("name")
                func_args = tool_json.get("arguments", {})

                yield accumulated_output + f"\n\n⚙️ `[Executing Tool: {func_name} | Args: {json.dumps(func_args)}]`...\n"
                
                # Run Tool
                tool_result = execute_tool(func_name, func_args, ip_folder, yolo_mode)
                
                # Feedback loop
                tool_feedback = f"\n<tool_response>\n{tool_result}\n</tool_response>\n"
                accumulated_output += tool_feedback
                yield accumulated_output
                
                chat_messages.append({"role": "user", "content": f"System Tool Output:\n{tool_result}\nAnalyze this and continue your task."})
            except Exception as e:
                error_fb = f"\n<tool_response>Error parsing JSON: {str(e)}</tool_response>\n"
                accumulated_output += error_fb
                chat_messages.append({"role": "user", "content": error_fb})
                yield accumulated_output
        else:
            # If no tools were called, the agent has finished responding.
            break

    with open(history_log_path, "a") as hf:
        hf.write(f"[{timestamp}] Assistant: {accumulated_output}\n")

def update_global_context(new_content):
    with open(global_context_path, "w") as f:
        f.write(new_content)
    return "✅ Global context updated successfully under /px!"

# =====================================================================
# 6. GRADIO WEB INTERFACE (CYBER-DARK STYLING)
# =====================================================================
custom_css = """
body { background-color: #0b0f19; color: #e2e8f0; font-family: 'Inter', sans-serif; }
.gradio-container { max-width: 1100px !important; margin: auto; padding-top: 20px; }
.gr-button-primary { background: linear-gradient(90deg, #3b82f6, #06b6d4) !important; border: none !important; }
.dark \.gr-panel { background-color: #111827 !important; border: 1px solid #1f2937 !important; }
code { color: #22d3ee !important; background-color: #1e293b !important; padding: 2px 6px; border-radius: 4px; }
pre code { color: #f8fafc !important; }
"""

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="cyan"), css=custom_css) as demo:
    gr.Markdown(
        """
        # ⚡ EVAL BETA 0.2 | PX TECH SOLUTIONS (AGENTIC ZEROGPU CORE)
        ### Autonomous Web Environment with ZeroGPU AI, Full Tool-Use, & YOLO God-Mode
        """
    )
    
    with gr.Row():
        gr.Markdown("🟢 **Status:** ZeroGPU Online | ⏱️ **Slot 1 Queue Active** | 🛠️ **Agent Loop:** Bash, Write, Read, Web Search")

    with gr.Accordion("⚙️ Engine Configuration, Global Context & Security", open=False):
        system_prompt_input = gr.Textbox(
            label="System Environment Prompt",
            value="You are EVAL BETA 0.2, an elite autonomous programming agent created by Pripro / PX TECH SOLUTIONS. Think logically and step-by-step.",
            lines=2
        )
        with gr.Row():
            temperature_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.3, step=0.05, label="Creativity / Temperature")
            yolo_checkbox = gr.Checkbox(label="🔥 Enable YOLO Mode (Disables Sandbox Security)", value=False)
            
        global_context_view = gr.Textbox(
            label=f"🌍 Shared Global Memory Workspace ({global_context_path})",
            value=open(global_context_path).read() if os.path.exists(global_context_path) else "",
            lines=4
        )
        update_btn = gr.Button("💾 Save Global Context Changes", variant="primary")
        update_status = gr.Textbox(label="Workspace Status", interactive=False)
        update_btn.click(update_global_context, inputs=[global_context_view], outputs=[update_status])

    gr.ChatInterface(
        fn=predict,
        additional_inputs=[system_prompt_input, temperature_slider, yolo_checkbox],
        textbox=gr.Textbox(placeholder="Instruct the AI: 'Search the web for Python 3.12 features', 'Write a script and run it via bash'...", container=False, scale=7)
    )
    
    gr.Markdown(
        """
        ---
        **Architecture & Security:** Enforces a strict queue (`concurrency_limit=1`) with Slot 1 allocation. All workspace data, logs, and IPs are securely sandboxed under the root `/px` architecture.
        """
    )

if __name__ == "__main__":
    print(f"🚀 Launching Agentic ZeroGPU chat server on port 7860...")
    demo.queue(default_concurrency_limit=1, max_size=20)
    demo.launch(server_name="0.0.0.0", server_port=7860)