"""
Nexa Package Manager — nexa install / list / remove / search
"""
from __future__ import annotations
import json
import os
import shutil
import urllib.request
import urllib.error
from pathlib import Path
from typing import Dict, List, Optional


# ── Package registry ──────────────────────────────────────────────────────────
# Built-in packages bundled with Nexa (resolved locally)
BUILTIN_REGISTRY: Dict[str, dict] = {
    "nexa-nlp": {
        "version": "1.0.0",
        "description": "Advanced NLP utilities: summarize, translate, ner",
        "entry": "nlp_utils.nexa",
        "builtin": True,
    },
    "nexa-vision": {
        "version": "1.0.0",
        "description": "Computer vision helpers: detect, segment, ocr",
        "entry": "vision_utils.nexa",
        "builtin": True,
    },
    "nexa-rl": {
        "version": "1.0.0",
        "description": "Reinforcement learning: environment, agent, train",
        "entry": "rl_utils.nexa",
        "builtin": True,
    },
    "nexa-data": {
        "version": "1.0.0",
        "description": "Data utilities: load_csv, split, describe, plot",
        "entry": "data_utils.nexa",
        "builtin": True,
    },
    "nexa-api": {
        "version": "1.0.0",
        "description": "HTTP client for APIs: get, post, auth",
        "entry": "api_utils.nexa",
        "builtin": True,
    },
}

# Community registry URL (can be overridden)
REGISTRY_URL = "https://raw.githubusercontent.com/nexa-lang/registry/main/packages.json"


class PackageManager:
    """Nexa package manager — installs packages into ~/.nexa/packages/"""

    def __init__(self, packages_dir: str = None):
        self.packages_dir = Path(packages_dir or os.path.expanduser("~/.nexa/packages"))
        self.lock_file = self.packages_dir.parent / "nexa.lock"
        self.packages_dir.mkdir(parents=True, exist_ok=True)

    # ── Lock file ─────────────────────────────────────────────────────────
    def _read_lock(self) -> dict:
        if self.lock_file.exists():
            return json.loads(self.lock_file.read_text())
        return {"installed": {}}

    def _write_lock(self, data: dict):
        self.lock_file.write_text(json.dumps(data, indent=2))

    # ── Registry ──────────────────────────────────────────────────────────
    def _fetch_registry(self) -> dict:
        """Fetch community registry; fall back to built-in on failure."""
        try:
            with urllib.request.urlopen(REGISTRY_URL, timeout=5) as r:
                return {**BUILTIN_REGISTRY, **json.loads(r.read())}
        except Exception:
            return BUILTIN_REGISTRY

    def _resolve(self, pkg_name: str) -> Optional[dict]:
        registry = self._fetch_registry()
        return registry.get(pkg_name)

    # ── Install ───────────────────────────────────────────────────────────
    def install(self, pkg_name: str) -> bool:
        lock = self._read_lock()
        
        info = None
        install_pkg_name = pkg_name
        
        if pkg_name.startswith("github:"):
            parts = pkg_name[7:].split("/")
            if len(parts) < 3:
                print("  ❌ Invalid github format. Use github:owner/repo/path/to/file.nexa")
                return False
            owner, repo = parts[0], parts[1]
            file_path = "/".join(parts[2:])
            url = f"https://raw.githubusercontent.com/{owner}/{repo}/main/{file_path}"
            entry_name = parts[-1]
            if not entry_name.endswith(".nexa"): entry_name += ".nexa"
            install_pkg_name = f"{owner}_{repo}_{entry_name[:-5]}"
            info = {
                "version": "latest",
                "description": f"GitHub package {owner}/{repo}/{file_path}",
                "entry": entry_name,
                "url": url
            }

        if install_pkg_name in lock.get("installed", {}):
            print(f"  ✅ {install_pkg_name} already installed "
                  f"(v{lock['installed'][install_pkg_name]['version']})")
            return True

        if not info:
            print(f"  🔍 Looking up '{pkg_name}'...")
            info = self._resolve(pkg_name)
            if info is None:
                print(f"  ❌ Package '{pkg_name}' not found in registry.")
                print(f"     Available: {', '.join(BUILTIN_REGISTRY)}")
                print(f"     Or use: github:owner/repo/path/file.nexa")
                return False

        pkg_dir = self.packages_dir / install_pkg_name
        pkg_dir.mkdir(exist_ok=True)

        if info.get("builtin"):
            # Generate a stub .nexa file for built-in packages
            self._create_builtin_stub(install_pkg_name, info, pkg_dir)
        else:
            # Download from URL
            url = info.get("url")
            if url:
                try:
                    dest = pkg_dir / info["entry"]
                    print(f"  ⬇️  Downloading from {url}...")
                    urllib.request.urlretrieve(url, dest)
                    
                    # Create __init__.py for Python package resolution
                    (pkg_dir / "__init__.py").write_text("")
                    
                    # Auto compile to python for Transpiler interop
                    try:
                        from .lexer import Lexer
                        from .parser import Parser
                        from .transpiler import Transpiler
                        src = dest.read_text(encoding="utf-8")
                        ast = Parser(Lexer(src).tokenize()).parse()
                        out_py = dest.with_suffix(".py")
                        out_py.write_text(Transpiler().transpile(ast), encoding="utf-8")
                    except Exception as e:
                        print(f"  ⚠️  Failed to precompile .py stub for {install_pkg_name}: {e}")
                except Exception as e:
                    print(f"  ❌ Download failed: {e}")
                    return False

        if "installed" not in lock:
            lock["installed"] = {}
        lock["installed"][install_pkg_name] = {
            "version": info["version"],
            "description": info["description"],
            "entry": str(pkg_dir / info["entry"]),
        }
        self._write_lock(lock)
        print(f"  ✅ Installed {install_pkg_name} v{info['version']}")
        print(f"     {info['description']}")
        if pkg_name.startswith("github:"):
            print(f"     Usage: from \"{install_pkg_name}\" import {{ ... }}")
        else:
            print(f"     Usage: import \"{pkg_dir / info['entry']}\" as {install_pkg_name.replace('-', '_')}")
        return True

    def _create_builtin_stub(self, pkg_name: str, info: dict, pkg_dir: Path):
        """Create a stub .nexa module for built-in packages."""
        stubs = {
            "nexa-nlp": """\
// nexa-nlp — Advanced NLP utilities
def summarize(text, max_len) {
    let words = nlp.tokenize(text)
    let n = len(words)
    if (n <= max_len) { return text }
    let result = ""
    let i = 0
    for (w in words) {
        if (i < max_len) { result = result + w + " " }
        i = i + 1
    }
    return result
}

def word_count(text) {
    return len(nlp.tokenize(text))
}

def analyze(text) {
    return {
        "sentiment": nlp.sentiment(text),
        "words": len(nlp.tokenize(text)),
        "embedding_dims": len(nlp.embed(text))
    }
}
""",
            "nexa-data": """\
// nexa-data — Data utilities
def describe(data) {
    let n = len(data)
    print("Records: " + str(n))
    return {"count": n}
}

def head(data, n) {
    let result = []
    let i = 0
    for (item in data) {
        if (i < n) { result = result + [item] }
        i = i + 1
    }
    return result
}

def filter_by(data, key, value) {
    let result = []
    for (item in data) {
        if (item[key] == value) {
            result = result + [item]
        }
    }
    return result
}
""",
            "nexa-rl": """\
// nexa-rl — Reinforcement Learning utilities
def create_agent(state_dim, action_dim) {
    return automl.neural_network(state_dim, action_dim, task="classification")
}

def epsilon_greedy(agent, state, epsilon) {
    // Simple epsilon-greedy policy
    let r = math.sqrt(1.0)  // placeholder random
    if (r < epsilon) {
        return 0  // random action
    }
    let actions = agent.predict([state])
    return actions[0]
}
""",
            "nexa-vision": """\
// nexa-vision — Vision utilities
def classify_batch(images, labels) {
    let results = []
    for (img in images) {
        let pred = vision.classify(img, labels)
        results = results + [pred]
    }
    return results
}

def preprocess(img, size) {
    return vision.resize(img, size, size)
}
""",
            "nexa-api": """\
// nexa-api — HTTP API client utilities  
def get_json(url) {
    print("Fetching: " + url)
    return {"url": url, "status": "ok"}
}

def post_json(url, data) {
    print("POST to: " + url)
    return {"url": url, "status": "ok", "data": data}
}
""",
        }
        content = stubs.get(pkg_name, f'// {pkg_name} stub\n')
        (pkg_dir / info["entry"]).write_text(content, encoding="utf-8")

    # ── List ──────────────────────────────────────────────────────────────
    def list_installed(self):
        lock = self._read_lock()
        installed = lock.get("installed", {})
        if not installed:
            print("  No packages installed. Run: nexa install <package>")
            return
        print(f"\n  📦 Installed Nexa packages ({len(installed)}):\n")
        for name, info in installed.items():
            print(f"  ✅ {name} v{info['version']}")
            print(f"     {info['description']}\n")

    # ── Search ────────────────────────────────────────────────────────────
    def search(self, query: str = ""):
        registry = self._fetch_registry()
        print(f"\n  🔍 Nexa Package Registry ({len(registry)} packages):\n")
        for name, info in registry.items():
            if not query or query.lower() in name.lower() or query.lower() in info["description"].lower():
                print(f"  📦 {name} v{info['version']}")
                print(f"     {info['description']}\n")

    # ── Remove ────────────────────────────────────────────────────────────
    def remove(self, pkg_name: str) -> bool:
        lock = self._read_lock()
        if pkg_name not in lock["installed"]:
            print(f"  ⚠️  '{pkg_name}' is not installed.")
            return False
        pkg_dir = self.packages_dir / pkg_name
        if pkg_dir.exists():
            shutil.rmtree(pkg_dir)
        del lock["installed"][pkg_name]
        self._write_lock(lock)
        print(f"  🗑️  Removed {pkg_name}")
        return True
