"""
Octopoda Multi-Tenant Manager
==============================
Each tenant (user account) gets their own isolated SQLite database.
No data leaks between tenants. Each has their own agents, memories,
knowledge graph, and dashboard view.

Architecture:
    ~/.synrix/data/tenants/{tenant_id}/synrix.db   — tenant's database
    ~/.synrix/data/tenants/{tenant_id}/backups/     — tenant's backups

Usage:
    manager = TenantManager.get_instance()
    backend = manager.get_backend("tenant_abc123")
    runtime = manager.get_runtime("tenant_abc123", "my_agent")
"""

import os
import time
import hashlib
import secrets
import threading
import json
import shutil
from typing import Dict, Optional, List
from dataclasses import dataclass, field


# ---------------------------------------------------------------------------
# Secure password hashing (PBKDF2 — stdlib, no extra deps)
# ---------------------------------------------------------------------------

_PBKDF2_ITERATIONS = 600_000  # OWASP 2023 recommendation for SHA256

def _hash_password(password: str) -> str:
    """Hash a password with PBKDF2-HMAC-SHA256 + random salt. Returns 'salt$hash'."""
    salt = secrets.token_hex(16)
    dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), _PBKDF2_ITERATIONS)
    return f"{salt}${dk.hex()}"

def _verify_password(password: str, stored: str) -> bool:
    """Verify a password against a stored 'salt$hash' string.
    Also accepts legacy plain SHA256 hashes for migration."""
    if "$" in stored:
        salt, hash_hex = stored.split("$", 1)
        dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), _PBKDF2_ITERATIONS)
        return secrets.compare_digest(dk.hex(), hash_hex)
    else:
        # Legacy SHA256 — verify and return True so caller can migrate
        legacy = hashlib.sha256(password.encode()).hexdigest()
        return secrets.compare_digest(legacy, stored)

def _is_legacy_hash(stored: str) -> bool:
    """Check if a stored hash is legacy (plain SHA256, no salt)."""
    return "$" not in stored


@dataclass
class TenantInfo:
    tenant_id: str
    email: str
    created_at: float
    plan: str = "free"
    max_agents: int = 10
    max_memories_per_agent: int = 10000
    api_key_hash: str = ""
    active: bool = True
    first_name: str = ""
    last_name: str = ""
    company: str = ""
    use_case: str = ""
    verified: bool = False


class TenantManager:
    """
    Manages isolated databases per tenant.
    Each tenant gets their own SQLite file and backend instance.
    """

    _instance = None
    _lock = threading.Lock()

    @classmethod
    def get_instance(cls, data_dir: str = None) -> "TenantManager":
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = cls(data_dir)
        return cls._instance

    def __init__(self, data_dir: str = None):
        self._data_dir = data_dir or os.path.expanduser("~/.synrix/data")
        self._tenants_dir = os.path.join(self._data_dir, "tenants")
        os.makedirs(self._tenants_dir, exist_ok=True)

        # Cache: tenant_id -> backend instance
        self._backends: Dict[str, object] = {}
        # Cache: (tenant_id, agent_id) -> AgentRuntime instance
        self._runtimes: Dict[tuple, object] = {}
        self._cache_lock = threading.Lock()

        # Master tenant registry (stored in root db)
        self._registry_path = os.path.join(self._data_dir, "tenant_registry.db")
        self._init_registry()

    def _init_registry(self):
        """Initialize the tenant registry database."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS tenants (
                    tenant_id TEXT PRIMARY KEY,
                    email TEXT UNIQUE NOT NULL,
                    password_hash TEXT NOT NULL,
                    created_at REAL NOT NULL,
                    plan TEXT DEFAULT 'free',
                    max_agents INTEGER DEFAULT 10,
                    max_memories_per_agent INTEGER DEFAULT 10000,
                    active INTEGER DEFAULT 1
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS api_keys (
                    key_hash TEXT PRIMARY KEY,
                    tenant_id TEXT NOT NULL,
                    key_prefix TEXT NOT NULL,
                    created_at REAL NOT NULL,
                    last_used REAL,
                    active INTEGER DEFAULT 1,
                    FOREIGN KEY (tenant_id) REFERENCES tenants(tenant_id)
                )
            """)
            conn.commit()

            # Add new columns for existing databases (safe — ALTER TABLE IF NOT EXISTS)
            for col, coltype in [
                ("first_name", "TEXT DEFAULT ''"),
                ("last_name", "TEXT DEFAULT ''"),
                ("company", "TEXT DEFAULT ''"),
                ("use_case", "TEXT DEFAULT ''"),
                ("verified", "INTEGER DEFAULT 0"),
            ]:
                try:
                    conn.execute(f"ALTER TABLE tenants ADD COLUMN {col} {coltype}")
                    conn.commit()
                except sqlite3.OperationalError:
                    pass  # Column already exists

            # Auto-verify all existing accounts (they signed up before verification existed)
            try:
                conn.execute("UPDATE tenants SET verified = 1 WHERE verified = 0 AND active = 1")
                conn.commit()
            except Exception:
                pass

    # ------------------------------------------------------------------
    # Tenant CRUD
    # ------------------------------------------------------------------

    def create_tenant(self, email: str, password: str, plan: str = "free",
                       first_name: str = "", last_name: str = "",
                       company: str = "", use_case: str = "") -> dict:
        """
        Create a new tenant account. Returns tenant_id + API key.
        The API key is shown once — we only store its hash.
        """
        import sqlite3

        tenant_id = hashlib.sha256(email.lower().encode()).hexdigest()[:16]
        password_hash = _hash_password(password)
        api_key = f"sk-octopoda-{secrets.token_urlsafe(32)}"
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
        key_prefix = api_key[:20]
        now = time.time()

        plan_limits = {
            "free": (100, 100000),
            "starter": (100, 100000),
            "pro": (500, 1000000),
            "enterprise": (10000, 10000000),
        }
        max_agents, max_memories = plan_limits.get(plan, (10, 10000))

        try:
            with sqlite3.connect(self._registry_path) as conn:
                conn.execute(
                    "INSERT INTO tenants (tenant_id, email, password_hash, created_at, plan, "
                    "max_agents, max_memories_per_agent, first_name, last_name, company, use_case, verified) "
                    "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 0)",
                    (tenant_id, email.lower(), password_hash, now, plan, max_agents, max_memories,
                     first_name, last_name, company, use_case),
                )
                conn.execute(
                    "INSERT INTO api_keys (key_hash, tenant_id, key_prefix, created_at, active) "
                    "VALUES (?, ?, ?, ?, 1)",
                    (key_hash, tenant_id, key_prefix, now),
                )
                conn.commit()
        except sqlite3.IntegrityError:
            return {"error": "Account already exists for this email", "success": False}

        # Create tenant's data directory
        tenant_dir = os.path.join(self._tenants_dir, tenant_id)
        os.makedirs(tenant_dir, exist_ok=True)
        os.makedirs(os.path.join(tenant_dir, "backups"), exist_ok=True)

        return {
            "success": True,
            "tenant_id": tenant_id,
            "api_key": api_key,
            "email": email.lower(),
            "first_name": first_name,
            "last_name": last_name,
            "plan": plan,
            "max_agents": max_agents,
            "max_memories_per_agent": max_memories,
            "warning": "Save your API key — it will not be shown again.",
        }

    def set_verified(self, email: str, verified: bool = True):
        """Mark an email as verified."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.execute(
                "UPDATE tenants SET verified = ? WHERE email = ? AND active = 1",
                (1 if verified else 0, email.lower()),
            )
            conn.commit()

    def get_tenant_by_email(self, email: str) -> Optional[dict]:
        """Get tenant info by email address."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT * FROM tenants WHERE email = ? AND active = 1",
                (email.lower(),),
            ).fetchone()
            return dict(row) if row else None

    def authenticate(self, email: str, password: str) -> Optional[dict]:
        """Authenticate with email + password. Returns tenant info or None."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT * FROM tenants WHERE email = ? AND active = 1",
                (email.lower(),),
            ).fetchone()
            if row and _verify_password(password, row["password_hash"]):
                # Auto-migrate legacy SHA256 hashes to PBKDF2
                if _is_legacy_hash(row["password_hash"]):
                    new_hash = _hash_password(password)
                    conn.execute(
                        "UPDATE tenants SET password_hash = ? WHERE tenant_id = ?",
                        (new_hash, row["tenant_id"]),
                    )
                    conn.commit()
                return dict(row)
        return None

    def verify_api_key(self, raw_key: str) -> Optional[dict]:
        """Verify an API key. Returns {tenant_id, plan, max_agents, ...} or None."""
        import sqlite3
        if not raw_key:
            return None
        if raw_key.startswith("Bearer "):
            raw_key = raw_key[7:]

        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT k.tenant_id, k.active as key_active, "
                "t.email, t.plan, t.max_agents, t.max_memories_per_agent, t.active as tenant_active, "
                "t.verified, t.first_name, t.last_name "
                "FROM api_keys k JOIN tenants t ON k.tenant_id = t.tenant_id "
                "WHERE k.key_hash = ?",
                (key_hash,),
            ).fetchone()
            if row and row["key_active"] and row["tenant_active"]:
                # Update last_used
                conn.execute(
                    "UPDATE api_keys SET last_used = ? WHERE key_hash = ?",
                    (time.time(), key_hash),
                )
                conn.commit()
                return dict(row)
        return None

    def create_session_key(self, tenant_id: str) -> Optional[str]:
        """Create an additional API key without deactivating existing ones (for login)."""
        import sqlite3
        api_key = f"sk-octopoda-{secrets.token_urlsafe(32)}"
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
        key_prefix = api_key[:20]
        now = time.time()
        with sqlite3.connect(self._registry_path) as conn:
            conn.execute(
                "INSERT INTO api_keys (key_hash, tenant_id, key_prefix, created_at, active) "
                "VALUES (?, ?, ?, ?, 1)",
                (key_hash, tenant_id, key_prefix, now),
            )
            conn.commit()
        return api_key

    def regenerate_api_key(self, tenant_id: str) -> Optional[str]:
        """Deactivate old keys, generate a new one. Returns raw key."""
        import sqlite3
        api_key = f"sk-octopoda-{secrets.token_urlsafe(32)}"
        key_hash = hashlib.sha256(api_key.encode()).hexdigest()
        key_prefix = api_key[:20]
        now = time.time()

        with sqlite3.connect(self._registry_path) as conn:
            # Deactivate old keys
            conn.execute(
                "UPDATE api_keys SET active = 0 WHERE tenant_id = ?",
                (tenant_id,),
            )
            # Insert new key
            conn.execute(
                "INSERT INTO api_keys (key_hash, tenant_id, key_prefix, created_at, active) "
                "VALUES (?, ?, ?, ?, 1)",
                (key_hash, tenant_id, key_prefix, now),
            )
            conn.commit()
        return api_key

    def get_tenant(self, tenant_id: str) -> Optional[dict]:
        """Get tenant info by ID."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT * FROM tenants WHERE tenant_id = ? AND active = 1",
                (tenant_id,),
            ).fetchone()
            return dict(row) if row else None

    def list_tenants(self) -> List[dict]:
        """List all active tenants (admin only)."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            rows = conn.execute(
                "SELECT tenant_id, email, plan, created_at, max_agents, max_memories_per_agent "
                "FROM tenants WHERE active = 1 ORDER BY created_at DESC"
            ).fetchall()
            return [dict(r) for r in rows]

    # ------------------------------------------------------------------
    # Backend / Runtime isolation
    # ------------------------------------------------------------------

    def _tenant_db_path(self, tenant_id: str) -> str:
        """Get the SQLite database path for a tenant."""
        return os.path.join(self._tenants_dir, tenant_id, "synrix.db")

    def get_backend(self, tenant_id: str):
        """Get or create an isolated backend for a tenant."""
        with self._cache_lock:
            if tenant_id in self._backends:
                return self._backends[tenant_id]

        from synrix.agent_backend import SynrixAgentBackend

        db_path = self._tenant_db_path(tenant_id)
        os.makedirs(os.path.dirname(db_path), exist_ok=True)

        backend = SynrixAgentBackend(backend="sqlite", sqlite_path=db_path)

        with self._cache_lock:
            self._backends[tenant_id] = backend
        return backend

    def get_runtime(self, tenant_id: str, agent_id: str):
        """Get or create an AgentRuntime isolated to a specific tenant."""
        cache_key = (tenant_id, agent_id)
        with self._cache_lock:
            if cache_key in self._runtimes:
                return self._runtimes[cache_key]

        # Check agent limit
        tenant = self.get_tenant(tenant_id)
        if tenant:
            agent_count = self.count_agents(tenant_id)
            if agent_count >= tenant["max_agents"]:
                raise TenantLimitError(
                    f"Agent limit reached: {agent_count}/{tenant['max_agents']} "
                    f"on {tenant['plan']} plan. Upgrade at https://octopoda.dev/pricing"
                )

        from synrix_runtime.api.runtime import AgentRuntime

        # Create a normal runtime (which initializes all internal state)
        # then swap its backend to the tenant-isolated one
        runtime = AgentRuntime(agent_id, agent_type="cloud", metadata={"tenant_id": tenant_id})

        # Override with tenant-isolated backend
        backend = self.get_backend(tenant_id)
        runtime.backend = backend

        # Also update the metrics collector to use tenant backend
        if hasattr(runtime, '_metrics') and runtime._metrics:
            runtime._metrics.backend = backend

        # Register agent in the tenant's DB so it shows up in list_agents
        now = time.time()
        try:
            backend.write(f"runtime:agents:{agent_id}:state", {"value": "running"})
            backend.write(f"runtime:agents:{agent_id}:type", {"value": "cloud"})
            backend.write(f"runtime:agents:{agent_id}:registered_at", {"value": now})
        except Exception:
            pass

        with self._cache_lock:
            self._runtimes[cache_key] = runtime
        return runtime

    def count_agents(self, tenant_id: str) -> int:
        """Count active agents for a tenant."""
        try:
            backend = self.get_backend(tenant_id)
            results = backend.query_prefix("runtime:agents:", limit=500)
            agents = set()
            for r in results:
                key = r.get("key", "")
                parts = key.split(":")
                if len(parts) >= 3 and parts[2] != "system":
                    agents.add(parts[2])
            return len(agents)
        except Exception:
            return 0

    def get_tenant_agents(self, tenant_id: str) -> List[dict]:
        """List all agents for a specific tenant."""
        try:
            backend = self.get_backend(tenant_id)
            results = backend.query_prefix("runtime:agents:", limit=500)
            agents = {}
            for r in results:
                key = r.get("key", "")
                parts = key.split(":")
                if len(parts) >= 3:
                    aid = parts[2]
                    if aid == "system":
                        continue
                    if aid not in agents:
                        agents[aid] = {"agent_id": aid}
                    if len(parts) > 3:
                        data = r.get("data", {})
                        value = data.get("value", data)
                        if isinstance(value, dict) and "value" in value:
                            value = value["value"]
                        agents[aid][parts[3]] = value
            return list(agents.values())
        except Exception:
            return []


    # ------------------------------------------------------------------
    # GDPR: Data export & account deletion
    # ------------------------------------------------------------------

    def export_tenant_data(self, tenant_id: str) -> dict:
        """Export all tenant data as a JSON-serializable dict (GDPR Article 20)."""
        tenant = self.get_tenant(tenant_id)
        if not tenant:
            return {"error": "Tenant not found"}

        # Remove internal fields
        export = {
            "account": {
                "tenant_id": tenant["tenant_id"],
                "email": tenant["email"],
                "plan": tenant["plan"],
                "created_at": tenant["created_at"],
                "max_agents": tenant["max_agents"],
            },
            "agents": [],
            "memories": [],
        }

        try:
            backend = self.get_backend(tenant_id)
            # Export agents
            agents = self.get_tenant_agents(tenant_id)
            export["agents"] = agents

            # Export all memories
            results = backend.query_prefix("agents:", limit=100000)
            for r in results:
                data = r.get("data", {})
                val = data.get("value", data)
                export["memories"].append({
                    "key": r.get("key", ""),
                    "value": val,
                })
        except Exception:
            pass

        return export

    def delete_tenant(self, tenant_id: str) -> dict:
        """Permanently delete a tenant account and all their data (GDPR Article 17)."""
        import sqlite3

        tenant = self.get_tenant(tenant_id)
        if not tenant:
            return {"error": "Tenant not found", "deleted": False}

        # 1. Remove from runtime caches
        with self._cache_lock:
            self._backends.pop(tenant_id, None)
            keys_to_remove = [k for k in self._runtimes if k[0] == tenant_id]
            for k in keys_to_remove:
                self._runtimes.pop(k, None)

        # 2. Delete tenant's data directory (SQLite DB + backups)
        tenant_dir = os.path.join(self._tenants_dir, tenant_id)
        if os.path.exists(tenant_dir):
            shutil.rmtree(tenant_dir, ignore_errors=True)

        # 3. Remove from registry
        with sqlite3.connect(self._registry_path) as conn:
            conn.execute("DELETE FROM api_keys WHERE tenant_id = ?", (tenant_id,))
            conn.execute("DELETE FROM tenants WHERE tenant_id = ?", (tenant_id,))
            conn.commit()

        return {"deleted": True, "tenant_id": tenant_id, "email": tenant["email"]}

    # ------------------------------------------------------------------
    # Usage stats
    # ------------------------------------------------------------------

    def get_tenant_usage(self, tenant_id: str) -> dict:
        """Get usage statistics for a tenant's dashboard."""
        tenant = self.get_tenant(tenant_id)
        if not tenant:
            return {"error": "Tenant not found"}

        agent_count = self.count_agents(tenant_id)
        memory_count = 0
        try:
            backend = self.get_backend(tenant_id)
            results = backend.query_prefix("agents:", limit=100000)
            memory_count = len(results)
        except Exception:
            pass

        return {
            "tenant_id": tenant_id,
            "plan": tenant["plan"],
            "agents": {"used": agent_count, "limit": tenant["max_agents"]},
            "memories": {"used": memory_count, "limit": tenant["max_memories_per_agent"] * max(agent_count, 1)},
            "created_at": tenant["created_at"],
        }

    # ------------------------------------------------------------------
    # Password management
    # ------------------------------------------------------------------

    def reset_password(self, email: str, new_password: str) -> dict:
        """Reset password without requiring old password (after code verification)."""
        import sqlite3
        if len(new_password) < 8:
            return {"error": "Password must be at least 8 characters", "success": False}
        new_hash = _hash_password(new_password)
        with sqlite3.connect(self._registry_path) as conn:
            result = conn.execute(
                "UPDATE tenants SET password_hash = ? WHERE email = ? AND active = 1",
                (new_hash, email.lower()),
            )
            conn.commit()
            if result.rowcount == 0:
                return {"error": "Account not found", "success": False}
        return {"success": True}

    def change_password(self, tenant_id: str, old_password: str, new_password: str) -> dict:
        """Change tenant password. Requires old password verification."""
        import sqlite3
        with sqlite3.connect(self._registry_path) as conn:
            conn.row_factory = sqlite3.Row
            row = conn.execute(
                "SELECT password_hash FROM tenants WHERE tenant_id = ? AND active = 1",
                (tenant_id,),
            ).fetchone()
            if not row:
                return {"error": "Tenant not found", "success": False}
            if not _verify_password(old_password, row["password_hash"]):
                return {"error": "Incorrect current password", "success": False}
            if len(new_password) < 8:
                return {"error": "Password must be at least 8 characters", "success": False}
            new_hash = _hash_password(new_password)
            conn.execute(
                "UPDATE tenants SET password_hash = ? WHERE tenant_id = ?",
                (new_hash, tenant_id),
            )
            conn.commit()
        return {"success": True}


class TenantLimitError(Exception):
    """Raised when a tenant exceeds their plan limits."""
    pass
