from __future__ import annotations

import hashlib
import json
import os
import secrets
import signal
import socket
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional

from harness_core import AuditLog, write_json_atomic
from launch_plan import validate_launch_plan


def canonical_sha(obj: Any) -> str:
    raw = json.dumps(
        obj,
        sort_keys=True,
        ensure_ascii=False,
        separators=(",", ":"),
    ).encode("utf-8")
    return hashlib.sha256(raw).hexdigest()


def proc_start_ticks(pid: int) -> int:
    # /proc/<pid>/stat field 22. comm can contain spaces inside parentheses,
    # so split only after the final ') '.
    text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
    tail = text[text.rfind(") ") + 2:].split()
    # Tail starts at field 3, so field 22 is tail index 19.
    return int(tail[19])


def proc_cmdline(pid: int) -> List[str]:
    raw = Path(f"/proc/{pid}/cmdline").read_bytes()
    return [
        x.decode("utf-8", errors="replace")
        for x in raw.split(b"\0")
        if x
    ]


def proc_state(pid: int):
    try:
        text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
    except FileNotFoundError:
        return None
    tail = text[text.rfind(") ") + 2:].split()
    if not tail:
        return None
    return tail[0]


def pid_exists(pid: int) -> bool:
    """
    Return True only for a live, non-zombie process.

    kill(pid, 0) alone is insufficient because a terminated child can remain
    briefly as a zombie until its parent reaps it. For lifecycle ownership,
    zombie means execution has ended and is therefore considered stopped.
    """
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False

    state = proc_state(pid)
    if state is None:
        return False
    if state == "Z":
        return False
    return True


def check_port_free(host: str, port: int) -> Dict[str, Any]:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0)
        s.bind((host, port))
        return {"status": "PASS", "free": True}
    except OSError as e:
        return {
            "status": "BLOCK",
            "free": False,
            "error": repr(e),
        }
    finally:
        s.close()


def choose_free_port(host: str = "127.0.0.1") -> int:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        s.bind((host, 0))
        return int(s.getsockname()[1])
    finally:
        s.close()


class RunOwnedRuntimeController:
    def __init__(
        self,
        *,
        run_dir: Path,
        audit: AuditLog,
        launch_plan: Dict[str, Any],
        readiness_url: Optional[str] = None,
        readiness_nonce: Optional[str] = None,
        startup_timeout: float = 15.0,
        stop_timeout: float = 10.0,
    ):
        self.run_dir = Path(run_dir)
        self.run_dir.mkdir(parents=True, exist_ok=True)
        self.audit = audit
        self.plan = dict(launch_plan)
        self.readiness_url = readiness_url
        self.readiness_nonce = readiness_nonce
        self.startup_timeout = float(startup_timeout)
        self.stop_timeout = float(stop_timeout)
        self.state_path = self.run_dir / "RUNTIME_STATE.json"
        self.stdout_path = self.run_dir / "runtime.stdout.log"
        self.stderr_path = self.run_dir / "runtime.stderr.log"

    def _read_state(self):
        if not self.state_path.is_file():
            return None
        return json.loads(
            self.state_path.read_text(encoding="utf-8")
        )

    def _write_state(self, state):
        write_json_atomic(self.state_path, state)

    def _clear_state(self):
        if self.state_path.exists():
            self.state_path.unlink()

    def validate_plan(self):
        result = validate_launch_plan(self.plan)
        if result["status"] != "PASS":
            node-7.example.invalid RuntimeError("LAUNCH_PLAN_BLOCK")
        return result

    def verify_ownership(self) -> Dict[str, Any]:
        state = self._read_state()
        if not state:
            return {
                "status": "BLOCK",
                "reason": "NO_RUNTIME_STATE",
            }

        pid = int(state["pid"])
        if not pid_exists(pid):
            return {
                "status": "BLOCK",
                "reason": "PID_NOT_ALIVE",
                "pid": pid,
            }

        actual_pgid = os.getpgid(pid)
        if actual_pgid != int(state["pgid"]):
            return {
                "status": "BLOCK",
                "reason": "PGID_MISMATCH",
                "expected": state["pgid"],
                "actual": actual_pgid,
            }

        actual_ticks = proc_start_ticks(pid)
        if actual_ticks != int(state["proc_start_ticks"]):
            return {
                "status": "BLOCK",
                "reason": "PROC_START_TICKS_MISMATCH",
                "expected": state["proc_start_ticks"],
                "actual": actual_ticks,
            }

        actual_cmdline = proc_cmdline(pid)
        expected_cmdline = state["expected_cmdline"]
        if actual_cmdline != expected_cmdline:
            return {
                "status": "BLOCK",
                "reason": "CMDLINE_MISMATCH",
                "expected": expected_cmdline,
                "actual": actual_cmdline,
            }

        if canonical_sha(actual_cmdline) != state["cmdline_sha256"]:
            return {
                "status": "BLOCK",
                "reason": "CMDLINE_SHA_MISMATCH",
            }

        return {
            "status": "PASS",
            "pid": pid,
            "pgid": actual_pgid,
            "proc_start_ticks": actual_ticks,
            "cmdline_sha256": state["cmdline_sha256"],
            "ownership_nonce": state["ownership_nonce"],
        }

    def probe(self) -> Dict[str, Any]:
        own = self.verify_ownership()
        if own["status"] != "PASS":
            return {
                "status": "BLOCK",
                "ownership": own,
            }

        if not self.readiness_url:
            return {
                "status": "PASS",
                "ownership": own,
                "readiness": "NOT_CONFIGURED",
            }

        try:
            with urllib.request.urlopen(
                self.readiness_url,
                timeout=2.0,
            ) as r:
                body = json.loads(
                    r.read().decode("utf-8")
                )
        except Exception as e:
            return {
                "status": "BLOCK",
                "ownership": own,
                "readiness": "HTTP_FAIL",
                "error": repr(e),
            }

        if self.readiness_nonce is not None:
            if body.get("nonce") != self.readiness_nonce:
                return {
                    "status": "BLOCK",
                    "ownership": own,
                    "readiness": "NONCE_MISMATCH",
                    "body": body,
                }

        return {
            "status": "PASS",
            "ownership": own,
            "readiness": "PASS",
            "body": body,
        }

    def start(self) -> Dict[str, Any]:
        plan_check = self.validate_plan()

        existing = self._read_state()
        if existing:
            own = self.verify_ownership()
            if own["status"] == "PASS":
                self.audit.append(
                    "RUNTIME_START_REUSE",
                    {"ownership": own},
                )
                return {
                    "status": "PASS",
                    "action": "ALREADY_RUNNING_RUN_OWNED",
                    "ownership": own,
                }
            # Fail closed on live-but-not-owned state. Clear only when PID is dead.
            pid = int(existing.get("pid", -1))
            if pid > 0 and pid_exists(pid):
                node-7.example.invalid RuntimeError(
                    "EXISTING_RUNTIME_STATE_NOT_OWNED:"
                    + json.dumps(own, ensure_ascii=False)
                )
            self.audit.append(
                "STALE_RUNTIME_STATE_CLEARED_BEFORE_START",
                {"old_state": existing, "ownership": own},
            )
            self._clear_state()

        executable = str(self.plan["executable"])
        argv = [str(x) for x in self.plan["argv"]]
        cmd = [executable] + argv
        env = dict(os.environ)
        for k, v in self.plan["env"].items():
            env[str(k)] = str(v)

        self.stdout_path.parent.mkdir(
            parents=True,
            exist_ok=True,
        )

        stdout_f = self.stdout_path.open("ab", buffering=0)
        stderr_f = self.stderr_path.open("ab", buffering=0)

        try:
            proc = subprocess.Popen(
                cmd,
                cwd=str(self.plan["cwd"]),
                env=env,
                stdin=subprocess.DEVNULL,
                stdout=stdout_f,
                stderr=stderr_f,
                shell=False,
                start_new_session=True,
                close_fds=True,
            )
        finally:
            stdout_f.close()
            stderr_f.close()

        pid = int(proc.pid)
        pgid = int(os.getpgid(pid))
        ticks = proc_start_ticks(pid)
        cmdline = proc_cmdline(pid)
        nonce = str(
            self.plan.get(
                "ownership_nonce",
                secrets.token_hex(16),
            )
        )

        state = {
            "schema_version": 1,
            "pid": pid,
            "pgid": pgid,
            "proc_start_ticks": ticks,
            "expected_cmdline": cmdline,
            "cmdline_sha256": canonical_sha(cmdline),
            "launch_plan_sha256": canonical_sha(self.plan),
            "ownership_nonce": nonce,
            "ownership_role": self.plan["ownership_role"],
            "started_unix": int(time.time()),
            "readiness_url": self.readiness_url,
            "stdout_path": str(self.stdout_path),
            "stderr_path": str(self.stderr_path),
        }
        self._write_state(state)

        self.audit.append(
            "RUNTIME_STARTED",
            {
                "pid": pid,
                "pgid": pgid,
                "proc_start_ticks": ticks,
                "cmdline_sha256": state["cmdline_sha256"],
                "launch_plan": plan_check,
            },
        )

        deadline = time.time() + self.startup_timeout
        last = None
        while time.time() < deadline:
            last = self.probe()
            if last["status"] == "PASS":
                self.audit.append(
                    "RUNTIME_READY",
                    last,
                )
                return {
                    "status": "PASS",
                    "action": "STARTED",
                    "probe": last,
                }
            if not pid_exists(pid):
                break
            time.sleep(0.1)

        self.audit.append(
            "RUNTIME_START_NOT_READY",
            {"last_probe": last},
        )
        # Do not broad-kill. Exact-stop only if ownership still verifies.
        own = self.verify_ownership()
        if own["status"] == "PASS":
            self.stop()
        node-7.example.invalid RuntimeError(
            "RUNTIME_READINESS_TIMEOUT:"
            + json.dumps(last, ensure_ascii=False)
        )

    def stop(self) -> Dict[str, Any]:
        state = self._read_state()
        if not state:
            return {
                "status": "PASS",
                "action": "NO_STATE",
            }

        own = self.verify_ownership()
        if own["status"] != "PASS":
            pid = int(state.get("pid", -1))
            if pid > 0 and not pid_exists(pid):
                self.audit.append(
                    "STALE_RUNTIME_STATE_CLEARED_ON_STOP",
                    {"state": state, "ownership": own},
                )
                self._clear_state()
                return {
                    "status": "PASS",
                    "action": "STALE_STATE_CLEARED",
                }
            node-7.example.invalid RuntimeError(
                "STOP_REFUSED_NOT_RUN_OWNED:"
                + json.dumps(own, ensure_ascii=False)
            )

        pgid = int(own["pgid"])
        pid = int(own["pid"])

        os.killpg(pgid, signal.SIGTERM)
        self.audit.append(
            "RUNTIME_SIGTERM_SENT",
            {
                "pid": pid,
                "pgid": pgid,
                "signal": "SIGTERM",
            },
        )

        deadline = time.time() + self.stop_timeout
        while time.time() < deadline:
            if not pid_exists(pid):
                self._clear_state()
                self.audit.append(
                    "RUNTIME_STOPPED",
                    {
                        "pid": pid,
                        "pgid": pgid,
                        "method": "SIGTERM",
                    },
                )
                return {
                    "status": "PASS",
                    "action": "SIGTERM_EXACT_RUN_OWNED_PGID",
                    "pid": pid,
                    "pgid": pgid,
                }
            time.sleep(0.1)

        # No SIGKILL escalation by design.
        self.audit.append(
            "RUNTIME_STOP_TIMEOUT",
            {
                "pid": pid,
                "pgid": pgid,
                "sigkill_escalation": False,
            },
        )
        return {
            "status": "BLOCK",
            "action": "SIGTERM_TIMEOUT_NO_SIGKILL",
            "pid": pid,
            "pgid": pgid,
        }

    def recover(self) -> Dict[str, Any]:
        state = self._read_state()
        if not state:
            return {
                "status": "PASS",
                "action": "NO_STATE",
            }

        own = self.verify_ownership()
        if own["status"] == "PASS":
            return {
                "status": "PASS",
                "action": "RUN_OWNED_PROCESS_STILL_ALIVE",
                "ownership": own,
            }

        pid = int(state.get("pid", -1))
        if pid > 0 and pid_exists(pid):
            return {
                "status": "BLOCK",
                "action": "LIVE_PROCESS_OWNERSHIP_MISMATCH",
                "ownership": own,
            }

        self.audit.append(
            "STALE_RUNTIME_STATE_RECOVERED",
            {
                "state": state,
                "ownership": own,
            },
        )
        self._clear_state()
        return {
            "status": "PASS",
            "action": "STALE_STATE_CLEARED",
        }