from __future__ import annotations

import logging
import os
import signal
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Optional, Set

logger = logging.getLogger("verlex.agent")

_DEFAULT_API_URL = "https://api.verlex.dev"


@dataclass
class ProcessInfo:
    pid: int
    name: str
    cmdline: list[str]
    script_path: Optional[str]
    cpu_percent: float
    memory_percent: float
    memory_mb: float
    username: str = ""
    create_time: float = 0.0

    @property
    def short_cmd(self) -> str:
        if self.script_path:
            return Path(self.script_path).name
        if self.cmdline:
            return " ".join(self.cmdline)[:60]
        return self.name


@dataclass
class OffloadRecord:
    pid: int
    script_path: str
    job_id: str
    submitted_at: float = field(default_factory=time.time)
    status: str = "submitted"


class ProcessScanner:
    def __init__(
        self,
        cpu_threshold: float = 85.0,
        mem_threshold: float = 85.0,
    ) -> None:
        self.cpu_threshold = cpu_threshold
        self.mem_threshold = mem_threshold
        self._own_pid = os.getpid()

    def scan(self) -> list[ProcessInfo]:
        import psutil

        heavy: list[ProcessInfo] = []

        for proc in psutil.process_iter(
            ["pid", "name", "cmdline", "cpu_percent", "memory_percent", "username", "create_time"]
        ):
            try:
                info = proc.info
                pid = info["pid"]

                if pid in (0, 1, self._own_pid):
                    continue

                if not self._is_python_process(info):
                    continue

                cpu = info.get("cpu_percent") or 0.0
                mem = info.get("memory_percent") or 0.0

                if cpu < self.cpu_threshold and mem < self.mem_threshold:
                    continue

                script = self._resolve_script(info.get("cmdline") or [])
                mem_mb = 0.0
                try:
                    mem_mb = proc.memory_info().rss / (1024 * 1024)
                except Exception:
                    pass

                heavy.append(
                    ProcessInfo(
                        pid=pid,
                        name=info.get("name", ""),
                        cmdline=info.get("cmdline") or [],
                        script_path=script,
                        cpu_percent=cpu,
                        memory_percent=mem,
                        memory_mb=mem_mb,
                        username=info.get("username") or "",
                        create_time=info.get("create_time") or 0.0,
                    )
                )
            except Exception:
                continue

        return heavy

    @staticmethod
    def _is_python_process(info: dict[str, Any]) -> bool:
        name = (info.get("name") or "").lower()
        if "python" in name:
            return True
        cmdline = info.get("cmdline") or []
        if cmdline and "python" in (cmdline[0] or "").lower():
            return True
        return False

    @staticmethod
    def _resolve_script(cmdline: list[str]) -> Optional[str]:
        for arg in cmdline:
            if not arg:
                continue
            if arg.startswith("-"):
                continue
            if arg.endswith(".py") and os.path.isfile(arg):
                return os.path.abspath(arg)
        return None


def _discover_local_modules(
    code: str,
    base_dir: str,
    _seen: Optional[Set[str]] = None,
) -> Dict[str, str]:
    import ast as _ast

    if _seen is None:
        _seen = set()

    try:
        tree = _ast.parse(code)
    except SyntaxError:
        return {}

    module_names: Set[str] = set()

    for node in _ast.walk(tree):
        if isinstance(node, _ast.ImportFrom) and node.module:
            module_names.add(node.module.split(".")[0])
        elif isinstance(node, _ast.Import):
            for alias in node.names:
                module_names.add(alias.name.split(".")[0])

    results: Dict[str, str] = {}

    for mod_name in sorted(module_names):
        if mod_name in _seen:
            continue

        candidate = os.path.join(base_dir, mod_name + ".py")
        if os.path.isfile(candidate):
            _seen.add(mod_name)
            mod_source = Path(candidate).read_text(encoding="utf-8", errors="replace")

            nested = _discover_local_modules(mod_source, base_dir, _seen)
            results.update(nested)
            results[mod_name] = mod_source
            continue

        pkg_init = os.path.join(base_dir, mod_name, "__init__.py")
        if os.path.isfile(pkg_init):
            _seen.add(mod_name)
            mod_source = Path(pkg_init).read_text(encoding="utf-8", errors="replace")

            nested = _discover_local_modules(mod_source, os.path.join(base_dir, mod_name), _seen)
            results.update(nested)
            results[mod_name] = mod_source

    return results


def _bundle_local_modules(code: str, script_path: str) -> str:
    base_dir = os.path.dirname(os.path.abspath(script_path))
    local_modules = _discover_local_modules(code, base_dir)

    if not local_modules:
        return code

    preamble_lines = [
        "# ---- verlex: injected local modules ----",
        "import sys as __vx_sys, types as __vx_types",
    ]

    for mod_name, mod_source in local_modules.items():
        preamble_lines.append(
            f"__vx_mod_{mod_name} = __vx_types.ModuleType({mod_name!r})\n"
            f"exec({mod_source!r}, __vx_mod_{mod_name}.__dict__)\n"
            f"__vx_sys.modules[{mod_name!r}] = __vx_mod_{mod_name}"
        )

    preamble_lines.append(
        "del __vx_sys, __vx_types"
        "\n# ---- verlex: end injected local modules ----\n"
    )

    return "\n".join(preamble_lines) + "\n" + code


def submit_code(
    code: str,
    file_path: Optional[str] = None,
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    fast: bool = False,
    gpu: Optional[str] = None,
    cpu: Optional[int] = None,
    memory: Optional[str] = None,
    priority: bool | None = None,
    flexible: bool | None = None,
) -> dict[str, Any]:
    import httpx

    if file_path:
        try:
            code = _bundle_local_modules(code, file_path)
        except Exception as exc:
            logger.warning("Failed to bundle local modules: %s", exc)

    api_key = (
        api_key
        or os.environ.get("VERLEX_API_KEY")
        or os.environ.get("GATEWAY_API_KEY")
    )
    if not api_key:
        raise RuntimeError(
            "No API key.  Pass api_key= or set VERLEX_API_KEY env var."
        )

    api_url = api_url or os.environ.get("VERLEX_API_URL", _DEFAULT_API_URL)

    resources: dict[str, Any] = {}
    if gpu:
        resources["gpu"] = gpu
    if cpu:
        resources["cpu"] = cpu
    if memory:
        resources["memory"] = memory

    payload = {
        "code": code,
        "file_path": file_path,
        "resources": resources,
        "fast": fast,
    }

    response = httpx.post(
        f"{api_url}/v1/jobs/submit/code",
        json=payload,
        headers={"Authorization": f"Bearer {api_key}"},
        timeout=60.0,
    )

    if response.status_code == 401:
        from verlex.errors import InvalidAPIKeyError
        raise InvalidAPIKeyError("Invalid API key.")
    if response.status_code == 402:
        from verlex.errors import InsufficientCreditsError
        data = response.json()
        raise InsufficientCreditsError(
            required=data.get("required", 0),
            available=data.get("available", 0),
        )
    if response.status_code == 429:
        from verlex.errors import RateLimitError
        raise RateLimitError(retry_after=None)
    if response.status_code != 200:
        from verlex.errors import ExecutionError
        raise ExecutionError(f"Submission failed ({response.status_code}): {response.text}")

    return response.json()


def poll_job(
    job_id: str,
    api_key: Optional[str] = None,
    api_url: Optional[str] = None,
    follow_logs: bool = True,
    poll_interval: float = 2.0,
) -> dict[str, Any]:
    import httpx

    api_key = (
        api_key
        or os.environ.get("VERLEX_API_KEY")
        or os.environ.get("GATEWAY_API_KEY")
    )
    api_url = api_url or os.environ.get("VERLEX_API_URL", _DEFAULT_API_URL)
    headers = {"Authorization": f"Bearer {api_key}"}

    log_offset = 0
    not_found_retries = 0

    while True:
        resp = httpx.get(
            f"{api_url}/v1/jobs/{job_id}",
            headers=headers,
            timeout=15.0,
        )
        if resp.status_code == 404:
            not_found_retries += 1
            if not_found_retries <= 5:
                time.sleep(min(not_found_retries * 2, 10))
                continue
            from verlex.errors import ExecutionError
            raise ExecutionError("Job not found after retries — server may have restarted.")
        if resp.status_code != 200:
            from verlex.errors import ExecutionError
            raise ExecutionError(f"Failed to check job ({resp.status_code}): {resp.text}")
        not_found_retries = 0

        data = resp.json()
        status = data.get("status", "unknown")

        if follow_logs:
            try:
                log_resp = httpx.get(
                    f"{api_url}/v1/jobs/{job_id}/output",
                    headers=headers,
                    params={"offset": log_offset},
                    timeout=10.0,
                )
                if log_resp.status_code == 200:
                    log_data = log_resp.json()
                    for entry in log_data.get("entries", []):
                        content = entry.get("content", entry.get("text", ""))
                        print(content, end="")
                    log_offset = log_data.get("current_index", log_offset)
            except Exception:
                pass

        if status in ("completed", "failed", "cancelled"):
            return data

        time.sleep(poll_interval)


class AgentDaemon:
    def __init__(
        self,
        api_key: Optional[str] = None,
        api_url: Optional[str] = None,
        fast: bool = False,
        threshold: float = 85.0,
        interval: float = 2.0,
        auto: bool = False,
        gpu: Optional[str] = None,
        cpu: Optional[int] = None,
        memory: Optional[str] = None,
    ) -> None:
        self.api_key = (
            api_key
            or os.environ.get("VERLEX_API_KEY")
            or os.environ.get("GATEWAY_API_KEY")
        )
        self.api_url = api_url or os.environ.get("VERLEX_API_URL", _DEFAULT_API_URL)
        self.fast = fast
        self.threshold = threshold
        self.interval = interval
        self.auto = auto
        self.gpu = gpu
        self.cpu = cpu
        self.memory = memory

        self.scanner = ProcessScanner(
            cpu_threshold=threshold,
            mem_threshold=threshold,
        )

        self._offloaded: Dict[int, OffloadRecord] = {}
        self._ignored_pids: Set[int] = set()
        self._running = False

    def watch(self) -> None:
        self._running = True

        prev_sigint = signal.getsignal(signal.SIGINT)

        def _stop(sig: int, frame: Any) -> None:
            self._running = False
            print("\nStopping agent watch...")

        signal.signal(signal.SIGINT, _stop)

        prev_sigterm = None
        if sys.platform != "win32":
            prev_sigterm = signal.getsignal(signal.SIGTERM)
            signal.signal(signal.SIGTERM, _stop)

        self._print_banner()

        import psutil
        psutil.cpu_percent(interval=None)
        for proc in psutil.process_iter(["cpu_percent"]):
            try:
                proc.cpu_percent()
            except Exception:
                pass
        time.sleep(1.0)

        try:
            while self._running:
                try:
                    self._tick()
                except Exception as exc:
                    logger.error("Scan tick failed: %s", exc)
                deadline = time.monotonic() + self.interval
                while self._running and time.monotonic() < deadline:
                    time.sleep(0.25)
        finally:
            signal.signal(signal.SIGINT, prev_sigint)
            if prev_sigterm is not None:
                signal.signal(signal.SIGTERM, prev_sigterm)

    def _tick(self) -> None:
        heavy = self.scanner.scan()

        for proc in heavy:
            if proc.pid in self._offloaded or proc.pid in self._ignored_pids:
                continue

            self._handle(proc)

    def _handle(self, proc: ProcessInfo) -> None:
        self._print_detection(proc)

        if not proc.script_path:
            print(f"    Cannot read source (no .py file on command line).  Skipping.\n")
            self._ignored_pids.add(proc.pid)
            return

        if not os.path.isfile(proc.script_path):
            print(f"    Script file not found on disk: {proc.script_path}.  Skipping.\n")
            self._ignored_pids.add(proc.pid)
            return

        if self.auto:
            self._offload(proc)
        else:
            self._prompt_offload(proc)

    def _prompt_offload(self, proc: ProcessInfo) -> None:
        try:
            answer = input("    Offload to cloud? [y/N/ignore] ").strip().lower()
        except EOFError:
            answer = "n"

        if answer in ("y", "yes"):
            self._offload(proc)
        elif answer in ("i", "ignore"):
            self._ignored_pids.add(proc.pid)
            print("    Ignored.\n")
        else:
            print("    Skipped.\n")

    def _offload(self, proc: ProcessInfo) -> None:
        try:
            code = Path(proc.script_path).read_text(encoding="utf-8", errors="replace")
        except Exception as exc:
            print(f"    Failed to read {proc.script_path}: {exc}\n")
            self._ignored_pids.add(proc.pid)
            return

        print(f"    Submitting {proc.short_cmd} to cloud...")

        try:
            result = submit_code(
                code=code,
                file_path=proc.script_path,
                api_key=self.api_key,
                api_url=self.api_url,
                fast=self.fast,
                gpu=self.gpu,
                cpu=self.cpu,
                memory=self.memory,
            )
        except Exception as exc:
            print(f"    Submission failed: {exc}\n")
            return

        job_id = result.get("job_id", "unknown")
        print(f"    Submitted!  job_id={job_id}")
        print(f"    Track with:  verlex status {job_id}")
        print(f"    Logs:        verlex logs {job_id} --follow\n")

        self._offloaded[proc.pid] = OffloadRecord(
            pid=proc.pid,
            script_path=proc.script_path or "",
            job_id=job_id,
        )

    def _print_banner(self) -> None:
        mode = "performance" if self.fast else "standard"
        auto_label = " (auto-offload)" if self.auto else ""

        print(f"Verlex Agent Watch{auto_label}")
        print(f"  Threshold : CPU or Memory > {self.threshold:.0f}%")
        print(f"  Mode      : {mode}")
        print(f"  Interval  : {self.interval:.1f}s")
        print(f"  Watching for heavy Python processes...  (Ctrl+C to stop)\n")

    @staticmethod
    def _print_detection(proc: ProcessInfo) -> None:
        print(f"  [{proc.pid}] {proc.short_cmd}")
        print(
            f"    CPU={proc.cpu_percent:.1f}%  "
            f"Mem={proc.memory_percent:.1f}% ({proc.memory_mb:.0f} MB)"
        )
        if proc.script_path:
            print(f"    Source: {proc.script_path}")
