#!/usr/bin/env python3
"""
ZODER V2.0 ULTIMATE AGENT (PUBLIC RELEASE)
Supports: Termux, Linux, macOS, Windows WSL
Cross-platform AI Terminal Agent with Tools & Encrypted Memory
"""

import os
import sys
import re
import json
import platform
import subprocess
import getpass
from datetime import datetime

# ============================================================
#  AUTO-INSTALL DEPENDENCIES
# ============================================================
def install_deps():
    deps = {
        "requests": "requests",
        "rich": "rich",
        "duckduckgo_search": "duckduckgo-search",
        "sqlcipher3": "sqlcipher3"
    }
    
    for module, pip_name in deps.items():
        try:
            __import__(module)
        except ImportError:
            print(f"⏳ Installing {pip_name}...")
            try:
                subprocess.check_call([sys.executable, "-m", "pip", "install", pip_name])
            except Exception as e:
                print(f"⚠️  Failed to install {pip_name}: {e}")
                if module == "sqlcipher3":
                    print("   Fallback to standard sqlite3 (UNENCRYPTED).")

install_deps()

import requests
from duckduckgo_search import DDGS

USE_RICH = False
try:
    from rich.console import Console
    from rich.panel import Panel
    from rich.prompt import Prompt, Confirm
    USE_RICH = True
except:
    pass

try:
    import sqlcipher3 as sqlite3_enc
    ENCRYPTED_DB = True
except:
    import sqlite3 as sqlite3_enc
    ENCRYPTED_DB = False

# ============================================================
#  CONFIGURATION
# ============================================================
OLLAMA_HOST = os.environ.get("ZODER_OLLAMA_HOST", "http://localhost:11434")
MODEL_NAME = os.environ.get("ZODER_MODEL_NAME", "zoder")
DB_PATH = os.path.expanduser("~/.zoder_memory.db")
SD_CPP_PATH = os.path.expanduser("~/llama.cpp/build/bin/llama-sd")
SD_MODEL_PATH = os.path.expanduser("~/models/sd-turbo-q4_0.gguf")

if USE_RICH:
    console = Console()


# ============================================================
#  UI HELPERS
# ============================================================
def print_ui(text, style="default"):
    if USE_RICH:
        if style == "banner":
            console.print(Panel.fit(text, border_style="cyan"))
        elif style == "user":
            console.print(f"[bold cyan]👤 You:[/bold cyan] {text}")
        elif style == "zoder":
            console.print(f"[bold magenta]🤖 Zoder:[/bold magenta] ", end="")
        elif style == "success":
            console.print(f"[bold green] ✅ {text}[/bold green]")
        elif style == "error":
            console.print(f"[bold red] ❌ {text}[/bold red]")
        elif style == "warning":
            console.print(f"[bold yellow] ⚠️  {text}[/bold yellow]")
        elif style == "info":
            console.print(f"[dim]{text}[/dim]")
        else:
            console.print(text)
    else:
        if style == "banner":
            print("\n" + "="*50 + f"\n{text}\n" + "="*50 + "\n")
        elif style == "user":
            print(f"\n👤 You: {text}")
        elif style == "zoder":
            print(f"\n🤖 Zoder: ", end="")
        elif style == "success":
            print(f"✅ {text}")
        elif style == "error":
            print(f"❌ {text}")
        elif style == "warning":
            print(f"⚠️  {text}")
        else:
            print(text)


# ============================================================
#  ENCRYPTED MEMORY (SQLCIPHER AES-256)
# ============================================================
class ZoderMemory:
    def __init__(self):
        self.conn = None
        self.cursor = None
        self.password = os.environ.get("ZODER_DB_PASS")
        
        if not self.password:
            if os.path.exists(DB_PATH):
                if USE_RICH:
                    self.password = Prompt.ask("[bold yellow]🔒 Enter Zoder Memory Password[/bold yellow]", password=True)
                else:
                    self.password = getpass.getpass("🔒 Enter Zoder Memory Password: ")
            else:
                if USE_RICH:
                    self.password = Prompt.ask("[bold green]🔑 Create New Memory Password[/bold green]", password=True)
                else:
                    self.password = getpass.getpass("🔑 Create New Memory Password: ")
        
        self._connect()

    def _connect(self):
        try:
            self.conn = sqlite3_enc.connect(DB_PATH)
            self.cursor = self.conn.cursor()
            
            if ENCRYPTED_DB:
                self.cursor.execute(f"PRAGMA key = '{self.password}';")
                self.cursor.execute("PRAGMA cipher_compatibility = 4;")
            
            self.cursor.execute("""
                CREATE TABLE IF NOT EXISTS memories (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp TEXT,
                    role TEXT,
                    content TEXT
                )
            """)
            self.conn.commit()
            
            self.cursor.execute("SELECT count(*) FROM memories")
            print_ui("Encrypted memory loaded successfully.", "success")
            
        except Exception as e:
            print_ui(f"Memory DB error (wrong password?): {e}", "error")
            sys.exit(1)

    def add(self, role, content):
        ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        self.cursor.execute("INSERT INTO memories (timestamp, role, content) VALUES (?, ?, ?)", (ts, role, content))
        self.conn.commit()

    def get_recent(self, limit=10):
        self.cursor.execute("SELECT role, content FROM memories ORDER BY id DESC LIMIT ?", (limit,))
        rows = self.cursor.fetchall()
        return [{"role": r[0], "content": r[1]} for r in reversed(rows)]


# ============================================================
#  TOOLS: WEB SEARCH, SHELL, IMAGE GEN
# ============================================================
def tool_web_search(query):
    print_ui(f"Searching web for: {query}...", "info")
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=3))
            if not results:
                return "No results found."
            
            summary = "Web Search Results:\n"
            for i, r in enumerate(results, 1):
                summary += f"{i}. {r['title']}\n   {r['body']}\n   Source: {r['href']}\n\n"
            return summary
    except Exception as e:
        return f"Web search failed: {e}"


def tool_shell_exec(command):
    if USE_RICH:
        if not Confirm.ask(f"[bold red]⚠️  Zoder wants to run:[/bold red] `{command}`\nAllow?"):
            return "Command execution denied by user."
    else:
        ans = input(f"⚠️  Zoder wants to run: {command}\nAllow? [y/N]: ").lower()
        if ans != 'y':
            return "Command execution denied by user."

    print_ui(f"Executing: {command}", "info")
    try:
        result = subprocess.run(command, shell=True, capture_output=True, text=True, timeout=30)
        output = result.stdout + result.stderr
        return output if output.strip() else "Command executed successfully (no output)."
    except Exception as e:
        return f"Shell execution failed: {e}"


def tool_image_gen(prompt):
    if not os.path.exists(SD_CPP_PATH) or not os.path.exists(SD_MODEL_PATH):
        return "Image generation tools not found. Please install llama.cpp with SD support and download sd-turbo-q4_0.gguf to ~/models/"
    
    output_file = os.path.join(os.getcwd(), f"zoder_img_{int(datetime.now().timestamp())}.png")
    cmd = f'{SD_CPP_PATH} -m "{SD_MODEL_PATH}" -p "{prompt}" -o "{output_file}" --steps 4'
    
    print_ui(f"Generating image (this takes a minute)...", "info")
    try:
        subprocess.run(cmd, shell=True, check=True, capture_output=True, timeout=120)
        return f"Image generated successfully at: {output_file}"
    except Exception as e:
        return f"Image generation failed: {e}"


# ============================================================
#  MAIN AGENT LOOP
# ============================================================
def clean_markdown(text):
    text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)
    text = re.sub(r'\*(.*?)\*', r'\1', text)
    text = re.sub(r'^#{1,6}\s+', '', text, flags=re.MULTILINE)
    text = re.sub(r'```[\s\S]*?```', '', text)
    text = re.sub(r'`(.*?)`', r'\1', text)
    text = re.sub(r'<[^>]+>', '', text)
    return text.strip()


def extract_and_execute_tools(text):
    patterns = {
        "WEB_SEARCH": r'\[WEB_SEARCH:\s*(.*?)\]',
        "SHELL_EXEC": r'\[SHELL_EXEC:\s*(.*?)\]',
        "IMAGE_GEN": r'\[IMAGE_GEN:\s*(.*?)\]',
        "CREATE_FILE": r'\[CREATE_FILE:\s*(.*?)\]\s*([\s\S]*?)\[/CREATE_FILE\]'
    }
    
    matches = re.findall(patterns["WEB_SEARCH"], text)
    for m in matches:
        result = tool_web_search(m.strip())
        text = text.replace(f"[WEB_SEARCH: {m}]", f"\n(Search Result: {result})")

    matches = re.findall(patterns["SHELL_EXEC"], text)
    for m in matches:
        result = tool_shell_exec(m.strip())
        text = text.replace(f"[SHELL_EXEC: {m}]", f"\n(Command Output: {result})")

    matches = re.findall(patterns["IMAGE_GEN"], text)
    for m in matches:
        result = tool_image_gen(m.strip())
        text = text.replace(f"[IMAGE_GEN: {m}]", f"\n({result})")

    matches = re.findall(patterns["CREATE_FILE"], text)
    for filename, content in matches:
        filepath = os.path.join(os.getcwd(), filename.strip())
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(content.strip())
            print_ui(f"File created: {filepath}", "success")
        except Exception as e:
            print_ui(f"Failed to create file: {e}", "error")
        text = re.sub(patterns["CREATE_FILE"], '', text)

    return text.strip()


def run_agent():
    os_info = "Android Termux" if "ANDROID_ROOT" in os.environ else f"{platform.system()} {platform.release()}"
    banner_text = (
        "[bold cyan]🤖 ZODER V2.0 ULTIMATE AGENT[/bold cyan]\n"
        f"[dim]Model: Zoder1.0-1B | System: {os_info}[/dim]\n"
        f"[dim]Memory: {'Encrypted (AES-256)' if ENCRYPTED_DB else 'Standard (Unencrypted)'}[/dim]"
    ) if USE_RICH else (
        f"🤖 ZODER V2.0 ULTIMATE AGENT\n"
        f"Model: Zoder1.0-1B | System: {os_info}\n"
        f"Memory: {'Encrypted (AES-256)' if ENCRYPTED_DB else 'Standard (Unencrypted)'}"
    )
    print_ui(banner_text, "banner")

    memory = ZoderMemory()
    print_ui("Type '/clear' to reset context, '/exit' to quit.", "info")

    while True:
        try:
            if USE_RICH:
                user_input = Prompt.ask("[bold cyan]👤 You[/bold cyan]")
            else:
                user_input = input("\n👤 You: ")
            
            if user_input.lower() == '/exit':
                print_ui("Goodbye!", "info")
                break
                
            if user_input.lower() == '/clear':
                print_ui("Context cleared for this session.", "success")
                continue

            if not user_input.strip():
                continue

            recent_mem = memory.get_recent(5)
            mem_context = "\n".join([f"{m['role']}: {m['content']}" for m in recent_mem])
            
            system_prompt = f"""You are Zoder, an advanced AI terminal agent. Model: Zoder1.0-1B.
RECENT MEMORY:
{mem_context}

TOOLS AVAILABLE (Use ONLY when needed, output the tag exactly):
1. To search web: [WEB_SEARCH: query]
2. To run shell command: [SHELL_EXEC: command]
3. To generate image: [IMAGE_GEN: prompt]
4. To create file: [CREATE_FILE: filename.ext]content[/CREATE_FILE]

RULES:
- NO MARKDOWN formatting in final response.
- Respond in the same language as user.
- If asked your name: "Zoder". If asked your model: "Zoder1.0-1B".
"""

            messages = [{"role": "system", "content": system_prompt}]
            messages.extend(recent_mem)
            messages.append({"role": "user", "content": user_input})

            print_ui("", "zoder")
            
            full_response = ""
            try:
                response = requests.post(
                    f"{OLLAMA_HOST}/api/chat",
                    json={"model": MODEL_NAME, "messages": messages, "stream": True},
                    stream=True, timeout=300
                )
                response.raise_for_status()
                
                for line in response.iter_lines():
                    if line:
                        try:
                            data = json.loads(line.decode('utf-8'))
                            if "message" in data and "content" in data["message"]:
                                token = data["message"]["content"]
                                full_response += token
                                sys.stdout.write(token)
                                sys.stdout.flush()
                        except:
                            continue
                print()
            except requests.exceptions.ConnectionError:
                print_ui("Cannot connect to Ollama! Is it running?", "error")
                continue

            processed = clean_markdown(full_response)
            final_response = extract_and_execute_tools(processed)
            
            memory.add("user", user_input)
            memory.add("assistant", final_response)
            
            print()

        except KeyboardInterrupt:
            print_ui("\nGoodbye!", "info")
            break
        except Exception as e:
            print_ui(f"Agent Error: {e}", "error")


if __name__ == "__main__":
    run_agent()