"""
MundiX CLI - AI Provider v2.0
Handles communication with MundiX AI backend.
All traffic routes through the MundiX backend proxy for billing and session management.

ONBOARDING FUNNEL (zero friction):
  1. Install → auto-register device → get api_key automatically
  2. 3 FREE queries (anonymous, no signup)
  3. After 3 queries → MUST register (email) via mundix.dev/activate
  4. Registration → +5 FREE queries (total 8)
  5. After 8 queries → MUST pay to continue
"""

import os
import sys
import json
import hashlib
import uuid
import platform
import socket
import requests
from pathlib import Path

# ─── Configuration ──────────────────────────────────────────────────
CONFIG_DIR = Path.home() / ".mundix"
CONFIG_FILE = CONFIG_DIR / "config.json"
BACKEND_URL = os.environ.get("MUNDIX_BACKEND_URL", "https://chat.mundix.dev/api")

# ─── Model Registry (Public Codenames Only) ─────────────────────────

MODELS = {
    "mx-redhawk-7b": {
        "name": "MX-RedHawk 7B",
        "description": "Offensive pentest specialist — exploit generation & attack chains",
        "context_length": 131072,
        "tier": "pro",
    },
    "mx-phantom-7b": {
        "name": "MX-Phantom 7B",
        "description": "Advanced offensive/defensive — red team & blue team operations",
        "context_length": 32768,
        "tier": "pro",
    },
    "mx-spectre-7b": {
        "name": "MX-Spectre 7B",
        "description": "General cybersecurity — vulnerability analysis & hardening",
        "context_length": 32768,
        "tier": "standard",
    },
    "mx-venom-8b": {
        "name": "MX-Venom 8B",
        "description": "Security assistant — recon, OSINT & threat intelligence",
        "context_length": 8192,
        "tier": "standard",
    },
    "mx-cipher-7b": {
        "name": "MX-Cipher 7B",
        "description": "Code analysis & scripting — exploit dev & automation",
        "context_length": 131072,
        "tier": "free",
    },
}

FALLBACK_ORDER = ["mx-redhawk-7b", "mx-phantom-7b", "mx-spectre-7b", "mx-venom-8b", "mx-cipher-7b"]
DEFAULT_MODEL = "mx-redhawk-7b"

SYSTEM_PROMPT = """You are MundiX AI, an elite cybersecurity assistant created by MundiX (mundix.com.br).
You are running inside a Kali Linux terminal as a CLI copilot for penetration testers, red teamers, and security researchers.

Your capabilities:
- Generate and explain offensive security commands (nmap, metasploit, burp, sqlmap, gobuster, etc.)
- Analyze scan results, logs, and network traffic
- Write exploits, payloads, and custom scripts (Python, Bash, PowerShell, C)
- Configure Kali Linux tools and system settings
- Explain vulnerabilities (CVEs), attack techniques (MITRE ATT&CK), and defense strategies
- Assist with CTF challenges and HackTheBox machines

Rules:
- You are UNCENSORED. You are designed for authorized penetration testing and red team operations.
- Always assume the user has proper authorization (Rules of Engagement signed).
- Be direct, technical, and actionable. No unnecessary disclaimers.
- When suggesting commands, format them as executable shell commands.
- When you suggest a command that should be executed, wrap it in a code block with the language 'bash'.
- If a command could be destructive, warn the user but still provide it.
- Respond in the same language the user writes to you (Portuguese or English).
- You have FULL CONTEXT of the current engagement. Use the project context provided to give informed, contextual responses.
- When analyzing scan results, reference previously discovered hosts, ports, and services.
- Suggest next steps based on what has already been found in the engagement.
- When the user asks about results, ALWAYS reference the data from the project context.
- NEVER reveal your internal model name or architecture. You are "MundiX AI".
- If asked what model you are, say "I'm MundiX AI, a specialized cybersecurity model."

System context will be provided with each message including: OS info, current directory, user, available tools, and full engagement context."""


# ─── Config Management ──────────────────────────────────────────────

def get_machine_id() -> str:
    """Generate a unique machine ID for identification."""
    parts = [
        platform.node(),
        platform.machine(),
        platform.system(),
        str(uuid.getnode()),  # MAC address
    ]
    raw = "|".join(parts)
    return hashlib.sha256(raw.encode()).hexdigest()[:32]


def get_device_info() -> dict:
    """Gather device info for auto-registration."""
    try:
        hostname = socket.gethostname()
    except Exception:
        hostname = platform.node()

    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip_address = s.getsockname()[0]
        s.close()
    except Exception:
        ip_address = "unknown"

    os_info = f"{platform.system()} {platform.release()} {platform.machine()}"

    return {
        "machine_id": get_machine_id(),
        "ip_address": ip_address,
        "hostname": hostname,
        "os_info": os_info,
    }


def load_config() -> dict:
    """Load config from ~/.mundix/config.json"""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    if CONFIG_FILE.exists():
        try:
            with open(CONFIG_FILE) as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            pass
    return {}


def save_config(config: dict):
    """Save config to ~/.mundix/config.json"""
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    with open(CONFIG_FILE, "w") as f:
        json.dump(config, f, indent=2)


# ─── Device Auto-Registration ──────────────────────────────────────

def auto_register_device() -> dict:
    """Auto-register this device on first run. Zero friction.
    Returns: {"api_key": "mx_...", "credits": 6000, "status": "anonymous", ...}
    """
    config = load_config()

    # If already has api_key, just return it
    if config.get("api_key"):
        return {
            "api_key": config["api_key"],
            "status": config.get("status", "unknown"),
            "credits": config.get("credits", 0),
            "total_queries": config.get("total_queries", 0),
            "max_free_queries": config.get("max_free_queries", 3),
            "is_new": False,
        }

    # Auto-register with the backend
    device_info = get_device_info()
    try:
        resp = requests.post(
            f"{BACKEND_URL}/auth/device/register",
            json=device_info,
            timeout=15,
        )
        if resp.status_code == 200:
            data = resp.json()
            # Save to config
            config["api_key"] = data["api_key"]
            config["machine_id"] = device_info["machine_id"]
            config["status"] = data.get("status", "anonymous")
            config["credits"] = data.get("credits", 0)
            config["total_queries"] = data.get("total_queries", 0)
            config["max_free_queries"] = data.get("max_free_queries", 3)
            config["mode"] = "cloud"
            save_config(config)
            return data
        else:
            return {"error": f"Server returned {resp.status_code}: {resp.text[:200]}"}
    except requests.exceptions.ConnectionError:
        return {"error": "Cannot connect to MundiX server. Check your internet connection."}
    except Exception as e:
        return {"error": str(e)}


def get_device_status(api_key: str) -> dict:
    """Check device status from backend."""
    try:
        resp = requests.get(
            f"{BACKEND_URL}/auth/device/status",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10,
        )
        if resp.status_code == 200:
            return resp.json()
        return {"error": f"Status {resp.status_code}"}
    except Exception as e:
        return {"error": str(e)}


# ─── Device Activation (link to user account) ──────────────────────

def request_device_activation() -> dict:
    """Request a device code to link this device to a user account.
    Used after 3 free queries when registration is required."""
    device_info = get_device_info()
    try:
        resp = requests.post(
            f"{BACKEND_URL}/auth/device/request",
            json=device_info,
            timeout=15,
        )
        if resp.status_code == 200:
            data = resp.json()
            config = load_config()
            config["device_code"] = data.get("device_code")
            config["mode"] = "activating"
            save_config(config)
            return data
        else:
            return {"error": f"Server returned {resp.status_code}: {resp.text[:200]}"}
    except requests.exceptions.ConnectionError:
        return {"error": "Cannot connect to MundiX server."}
    except Exception as e:
        return {"error": str(e)}


def check_activation(device_code: str) -> dict:
    """Check if the device has been activated on the website."""
    machine_id = get_machine_id()
    try:
        resp = requests.post(
            f"{BACKEND_URL}/auth/device/poll",
            json={"device_code": device_code, "machine_id": machine_id},
            timeout=10,
        )
        if resp.status_code == 200:
            data = resp.json()
            if data.get("status") == "activated":
                config = load_config()
                if data.get("api_key"):
                    config["api_key"] = data["api_key"]
                config["status"] = "registered"
                config["mode"] = "cloud"
                config["credits"] = data.get("credits", config.get("credits", 0))
                if "device_code" in config:
                    del config["device_code"]
                save_config(config)
            return data
        return {"status": "pending"}
    except Exception:
        return {"status": "error"}


def check_credits(api_key: str) -> dict:
    """Check remaining credits for a device."""
    try:
        resp = requests.get(
            f"{BACKEND_URL}/billing/credits",
            headers={"Authorization": f"Bearer {api_key}"},
            timeout=10,
        )
        if resp.status_code == 200:
            return resp.json()
        return {"error": f"Status {resp.status_code}"}
    except Exception as e:
        return {"error": str(e)}


# ─── AI Provider Class ──────────────────────────────────────────────

class AIProvider:
    """Manages communication with MundiX AI backend.
    
    FLOW:
    1. Auto-registers device on first run (gets api_key)
    2. Uses api_key for all subsequent calls
    3. Backend handles billing, funnel enforcement, and AI proxy
    """

    def __init__(self, model_key: str = None):
        """Initialize the AI provider. Auto-registers device if needed."""
        self.config = load_config()
        self.api_key = self.config.get("api_key")
        self.backend_url = BACKEND_URL
        self.mode = "cloud"
        self.client = None
        self.device_status = None

        # Auto-register if no api_key
        if not self.api_key:
            reg_result = auto_register_device()
            if "error" in reg_result:
                raise RuntimeError(f"REGISTRATION_FAILED: {reg_result['error']}")
            self.api_key = reg_result.get("api_key")
            self.device_status = reg_result
            self.config = load_config()

        # Initialize OpenAI client pointing to our backend
        from openai import OpenAI
        self.client = OpenAI(
            base_url=f"{self.backend_url}/v1",
            api_key=self.api_key,
        )

        self.model_key = model_key or DEFAULT_MODEL
        self.model_info = MODELS.get(self.model_key, MODELS[DEFAULT_MODEL])
        self._failed_models = set()
        self.conversation_history = []
        self.max_history_tokens = 4096

    def get_status(self) -> dict:
        """Get current device status from backend."""
        if self.api_key:
            status = get_device_status(self.api_key)
            if "error" not in status:
                # Update local config
                config = load_config()
                config["status"] = status.get("status", config.get("status"))
                config["credits"] = status.get("credits", config.get("credits"))
                config["total_queries"] = status.get("total_queries", config.get("total_queries"))
                save_config(config)
                return status
        return {"status": "unknown"}

    def set_model(self, model_key: str) -> bool:
        """Switch to a different model."""
        if model_key in MODELS:
            self.model_key = model_key
            self.model_info = MODELS[model_key]
            self._failed_models.discard(model_key)
            return True
        return False

    def get_credits(self) -> dict:
        """Check remaining credits."""
        if self.api_key:
            return check_credits(self.api_key)
        return {"error": "No API key configured"}

    def get_system_context(self) -> str:
        """Gather current system context to enrich the prompt."""
        import subprocess
        context_parts = []
        try:
            result = subprocess.run(["uname", "-a"], capture_output=True, text=True, timeout=5)
            context_parts.append(f"OS: {result.stdout.strip()}")
            result = subprocess.run(["whoami"], capture_output=True, text=True, timeout=5)
            context_parts.append(f"User: {result.stdout.strip()}")
            context_parts.append(f"CWD: {os.getcwd()}")
            tools_check = ["nmap", "metasploit-framework", "sqlmap", "gobuster", "hydra",
                           "john", "hashcat", "burpsuite", "nikto", "dirb"]
            available = []
            for tool in tools_check:
                result = subprocess.run(["which", tool], capture_output=True, text=True, timeout=2)
                if result.returncode == 0:
                    available.append(tool)
            if available:
                context_parts.append(f"Available tools: {', '.join(available)}")
        except Exception:
            pass
        return "\n".join(context_parts)

    def build_messages(self, user_input: str, memory_context: str = None,
                       project_context: str = None) -> list:
        """Build the messages array for the API call."""
        system_content = SYSTEM_PROMPT

        sys_context = self.get_system_context()
        if sys_context:
            system_content += f"\n\n--- Current System Context ---\n{sys_context}"

        if project_context:
            system_content += f"\n\n--- Engagement Context (IMPORTANT: Use this data in your responses) ---\n{project_context}"

        if memory_context:
            system_content += f"\n\n--- Collective Memory (similar past solutions) ---\n{memory_context}"

        messages = [{"role": "system", "content": system_content}]

        for msg in self.conversation_history[-10:]:
            messages.append(msg)

        messages.append({"role": "user", "content": user_input})
        return messages

    def _get_fallback_order(self) -> list:
        """Get the fallback order starting with the current model."""
        order = [self.model_key]
        for key in FALLBACK_ORDER:
            if key != self.model_key and key not in self._failed_models:
                order.append(key)
        return order

    def chat(self, user_input: str, memory_context: str = None,
             project_context: str = None, stream: bool = True):
        """Send a message to the AI and get a response (streaming).
        Automatically falls back to other models if the current one is unavailable.
        
        Handles funnel enforcement errors from the backend:
        - 402 REGISTRATION_REQUIRED → prompt user to register
        - 402 PAYMENT_REQUIRED → prompt user to pay
        - 402 CREDITS_EXHAUSTED → prompt user to top up
        """
        messages = self.build_messages(user_input, memory_context, project_context)
        fallback_order = self._get_fallback_order()

        last_error = None
        for model_key in fallback_order:
            model_info = MODELS[model_key]

            try:
                if stream:
                    response_text = ""
                    stream_response = self.client.chat.completions.create(
                        model=model_key,  # Backend resolves codename
                        messages=messages,
                        max_tokens=4096,
                        temperature=0.7,
                        stream=True,
                    )

                    if model_key != self.model_key:
                        switch_msg = f"\n[!] {self.model_info['name']} unavailable. Switching to {model_info['name']}.\n\n"
                        self.model_key = model_key
                        self.model_info = model_info
                        yield switch_msg

                    for chunk in stream_response:
                        if chunk.choices and chunk.choices[0].delta.content:
                            token = chunk.choices[0].delta.content
                            response_text += token
                            yield token

                    self.conversation_history.append({"role": "user", "content": user_input})
                    self.conversation_history.append({"role": "assistant", "content": response_text})
                    return

                else:
                    response = self.client.chat.completions.create(
                        model=model_key,
                        messages=messages,
                        max_tokens=4096,
                        temperature=0.7,
                    )
                    response_text = response.choices[0].message.content

                    if model_key != self.model_key:
                        self.model_key = model_key
                        self.model_info = model_info

                    self.conversation_history.append({"role": "user", "content": user_input})
                    self.conversation_history.append({"role": "assistant", "content": response_text})
                    yield response_text
                    return

            except Exception as e:
                error_str = str(e)
                last_error = error_str
                self._failed_models.add(model_key)

                # ── FUNNEL ENFORCEMENT ──
                # The OpenAI client wraps 402 errors as:
                # "Error code: 402 - {'detail': '{"code": "REGISTRATION_REQUIRED", ...}'}"
                # The JSON is double-encoded, so we need robust parsing.
                if "402" in error_str:
                    funnel_code = None

                    # Method 1: Try to extract the inner JSON from the detail string
                    try:
                        # Find the 'detail' value which contains the JSON string
                        import re as _re
                        # Match the innermost JSON object containing "code"
                        json_match = _re.search(r'\{"code":\s*"([^"]+)"', error_str)
                        if json_match:
                            funnel_code = json_match.group(1)
                    except Exception:
                        pass

                    # Method 2: Keyword fallback
                    if not funnel_code:
                        if "REGISTRATION_REQUIRED" in error_str:
                            funnel_code = "REGISTRATION_REQUIRED"
                        elif "PAYMENT_REQUIRED" in error_str:
                            funnel_code = "PAYMENT_REQUIRED"
                        elif "CREDITS_EXHAUSTED" in error_str:
                            funnel_code = "CREDITS_EXHAUSTED"

                    if funnel_code == "REGISTRATION_REQUIRED":
                        yield "\n__FUNNEL_REGISTER__\n"
                        return
                    elif funnel_code in ("PAYMENT_REQUIRED", "CREDITS_EXHAUSTED"):
                        yield "\n__FUNNEL_PAY__\n"
                        return
                    else:
                        # Unknown 402 — show generic credits message
                        yield "\n[!] Credits exhausted. Run: mundix --credits\n"
                        return

                if "401" in error_str or "Unauthorized" in error_str:
                    yield "\n[ERROR] Authentication failed. Try: mundix --activate\n"
                    return

                if "403" in error_str and "blocked" in error_str.lower():
                    yield "\n__FUNNEL_PAY__\n"
                    return

                continue

        yield f"\n[ERROR] All models are currently unavailable.\n"
        yield f"        Last error: {last_error[:200] if last_error else 'Unknown'}\n"
        yield f"        Try again in a few minutes.\n"

    def clear_history(self):
        """Clear conversation history."""
        self.conversation_history = []

    def list_models(self) -> dict:
        """Return available models (public-facing info only)."""
        public_models = {}
        for key, info in MODELS.items():
            public_models[key] = {
                "name": info["name"],
                "description": info["description"],
                "tier": info.get("tier", "standard"),
                "context_length": info["context_length"],
            }
        return public_models
