"""In-process capture store for the Mugi browser extension.

Manual extension clicks still use the compatibility ``_LATEST_CAPTURE`` slot.
Agent-requested captures use single-flight jobs keyed by ``job_id`` so content
from one chat thread cannot be consumed by another.
"""
from __future__ import annotations

import copy
import ipaddress
import json
import os
import threading
import time
import uuid
from datetime import datetime, timezone
from typing import Any
from urllib.parse import urlsplit, urlunsplit

_LOCK = threading.RLock()
_CONDITION = threading.Condition(_LOCK)
_LATEST_CAPTURE: dict[str, Any] | None = None
_JOBS: dict[str, dict[str, Any]] = {}
_AGENT_CAPTURE_ENABLED = False
_CREATE_JOB_HITS: list[float] = []

PENDING_TTL_SECONDS = 10 * 60
CLAIMED_TTL_SECONDS = 5 * 60
TERMINAL_TTL_SECONDS = 30 * 60
CREATE_JOB_RATE_LIMIT_PER_MINUTE = 6
MAX_URL_LENGTH = 2048
AUDIT_LOG_PATH = os.path.expanduser("~/.mugi/browser_extension_audit.jsonl")

STATUS_PENDING = "pending"
STATUS_CLAIMED = "claimed"
STATUS_COMPLETE = "complete"
STATUS_FAILED = "failed"
STATUS_EXPIRED = "expired"

_BIDI_CONTROL_CHARS = frozenset(
    {
        "\u202a",
        "\u202b",
        "\u202c",
        "\u202d",
        "\u202e",
        "\u2066",
        "\u2067",
        "\u2068",
        "\u2069",
    }
)

SENSITIVE_HOST_SEEDS = frozenset(
    {
        "accounts.google.com",
        "auth0.com",
        "bankofamerica.com",
        "bitwarden.com",
        "chase.com",
        "fidelity.com",
        "gmail.com",
        "healthcare.gov",
        "icloud.com",
        "irs.gov",
        "lastpass.com",
        "login.gov",
        "mail.google.com",
        "okta.com",
        "paypal.com",
        "schwab.com",
        "vanguard.com",
        "wellsfargo.com",
        "1password.com",
    }
)


class BrowserCaptureError(Exception):
    """Structured error for agent-requested browser capture state."""

    def __init__(self, code: str, message: str, *, status_code: int = 400):
        super().__init__(message)
        self.code = code
        self.message = message
        self.status_code = status_code


def _now_iso() -> str:
    return datetime.now().isoformat(timespec="seconds")


def _audit_ts() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="seconds")


def _registrable_host(host: str) -> str:
    labels = [p for p in host.strip(".").lower().split(".") if p]
    if len(labels) >= 2:
        return ".".join(labels[-2:])
    return labels[0] if labels else ""


def _host_matches_seed(host: str, seed: str) -> bool:
    h = host.strip(".").lower()
    s = seed.strip(".").lower()
    return h == s or h.endswith("." + s)


def is_sensitive_host(host: str) -> bool:
    """Return True when *host* is in the built-in sensitive seed list."""
    return any(_host_matches_seed(host, seed) for seed in SENSITIVE_HOST_SEEDS)


def validate_capture_url(raw_url: str, *, allow_punycode: bool = False) -> dict[str, Any]:
    """Validate a URL for agent-requested capture and return normalized fields."""
    raw = str(raw_url or "").strip()
    if not raw:
        raise BrowserCaptureError("invalid_url", "Missing URL.")
    if len(raw) > MAX_URL_LENGTH:
        raise BrowserCaptureError("url_too_long", "URL is too long.")
    if any(ord(ch) < 32 for ch in raw) or any(ch in _BIDI_CONTROL_CHARS for ch in raw):
        raise BrowserCaptureError("invalid_url", "URL contains unsafe control characters.")
    try:
        parsed = urlsplit(raw)
    except Exception as exc:
        raise BrowserCaptureError("invalid_url", "URL could not be parsed.") from exc
    scheme = parsed.scheme.lower()
    if scheme not in ("http", "https"):
        raise BrowserCaptureError("unsupported_scheme", "Only http(s) URLs can be captured.")
    if parsed.username or parsed.password:
        raise BrowserCaptureError("userinfo_not_allowed", "URLs with embedded credentials are not allowed.")
    host = (parsed.hostname or "").strip(".").lower()
    if not host:
        raise BrowserCaptureError("invalid_host", "URL must include a host.")
    if host.startswith("xn--") or ".xn--" in host:
        if not allow_punycode:
            raise BrowserCaptureError("punycode_host", "Punycode/IDN hosts require manual capture.")
    try:
        ip = ipaddress.ip_address(host)
        if ip.is_link_local or ip.is_multicast or ip.is_unspecified:
            raise BrowserCaptureError("unsafe_ip", "Unsafe IP address is not allowed.")
    except ValueError:
        pass
    try:
        parsed_port = parsed.port
    except ValueError as exc:
        raise BrowserCaptureError("invalid_port", "URL port is invalid.") from exc
    port = f":{parsed_port}" if parsed_port else ""
    netloc = f"{host}{port}"
    normalized = urlunsplit((scheme, netloc, parsed.path or "/", parsed.query, parsed.fragment))
    return {
        "url": normalized,
        "scheme": scheme,
        "host": host,
        "display_host": _registrable_host(host),
        "sensitive": is_sensitive_host(host),
        "punycode": host.startswith("xn--") or ".xn--" in host,
    }


def set_agent_capture_enabled(enabled: bool) -> bool:
    """Set the in-process agent capture killswitch."""
    global _AGENT_CAPTURE_ENABLED
    with _CONDITION:
        _AGENT_CAPTURE_ENABLED = bool(enabled)
        _CONDITION.notify_all()
        return _AGENT_CAPTURE_ENABLED


def agent_capture_enabled() -> bool:
    """Return whether agent-requested capture is enabled."""
    with _LOCK:
        return bool(_AGENT_CAPTURE_ENABLED)


def _append_audit(job: dict[str, Any], status: str, *, decision: str = "", error: str = "") -> None:
    entry = {
        "ts": _audit_ts(),
        "job_id": job.get("job_id") or "",
        "thread_id": job.get("thread_id") or "",
        "host": job.get("host") or "",
        "status": status,
        "decision": decision or job.get("decision") or "",
        "error": error or "",
    }
    try:
        os.makedirs(os.path.dirname(AUDIT_LOG_PATH), exist_ok=True)
        fd = os.open(AUDIT_LOG_PATH, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        try:
            os.write(fd, (json.dumps(entry, ensure_ascii=False) + "\n").encode("utf-8"))
        finally:
            os.close(fd)
        try:
            os.chmod(AUDIT_LOG_PATH, 0o600)
        except OSError:
            pass
    except OSError:
        pass


def _terminal(status: str) -> bool:
    return status in {STATUS_COMPLETE, STATUS_FAILED, STATUS_EXPIRED}


def _expire_and_prune_locked(now: float | None = None) -> None:
    now = time.time() if now is None else now
    terminal_cutoff = now - TERMINAL_TTL_SECONDS
    for job_id, job in list(_JOBS.items()):
        status = str(job.get("status") or "")
        created = float(job.get("created_monotonic") or now)
        claimed = float(job.get("claimed_monotonic") or created)
        terminal_at = float(job.get("terminal_monotonic") or 0)
        if status == STATUS_PENDING and now - created > PENDING_TTL_SECONDS:
            job["status"] = STATUS_EXPIRED
            job["terminal_at"] = _now_iso()
            job["terminal_monotonic"] = now
            job["error"] = "expired"
            _append_audit(job, STATUS_EXPIRED, decision="timeout", error="pending_ttl")
        elif status == STATUS_CLAIMED and now - claimed > CLAIMED_TTL_SECONDS:
            job["status"] = STATUS_EXPIRED
            job["terminal_at"] = _now_iso()
            job["terminal_monotonic"] = now
            job["error"] = "expired"
            _append_audit(job, STATUS_EXPIRED, decision="timeout", error="claimed_ttl")
        elif _terminal(status) and terminal_at and terminal_at < terminal_cutoff:
            _JOBS.pop(job_id, None)


def _rate_limited_locked(now: float) -> bool:
    cutoff = now - 60
    while _CREATE_JOB_HITS and _CREATE_JOB_HITS[0] < cutoff:
        _CREATE_JOB_HITS.pop(0)
    if len(_CREATE_JOB_HITS) >= CREATE_JOB_RATE_LIMIT_PER_MINUTE:
        return True
    _CREATE_JOB_HITS.append(now)
    return False


def create_job(url: str, *, thread_id: str | None = None, now: float | None = None) -> dict[str, Any]:
    """Create a single-flight agent-requested capture job."""
    info = validate_capture_url(url)
    if info["sensitive"]:
        raise BrowserCaptureError("blocked_sensitive", "Sensitive hosts require manual capture.", status_code=403)
    with _CONDITION:
        if not _AGENT_CAPTURE_ENABLED:
            raise BrowserCaptureError("capture_disabled", "Agent browser capture is disabled.", status_code=503)
        now_mono = time.time() if now is None else now
        _expire_and_prune_locked(now_mono)
        if _rate_limited_locked(now_mono):
            raise BrowserCaptureError("rate_limited", "Too many browser capture requests.", status_code=429)
        for job in _JOBS.values():
            if job.get("status") in (STATUS_PENDING, STATUS_CLAIMED):
                raise BrowserCaptureError("busy", "A browser capture request is already in progress.", status_code=409)
        job_id = uuid.uuid4().hex
        job = {
            "job_id": job_id,
            "thread_id": str(thread_id or "default")[:128],
            "url": info["url"],
            "requested_url": info["url"],
            "host": info["display_host"],
            "full_host": info["host"],
            "status": STATUS_PENDING,
            "created_at": _now_iso(),
            "created_monotonic": now_mono,
            "punycode": bool(info.get("punycode")),
            "result": None,
            "error": "",
        }
        _JOBS[job_id] = job
        _append_audit(job, STATUS_PENDING, decision="created")
        _CONDITION.notify_all()
        return _public_job(job, include_url=True)


def _public_job(job: dict[str, Any], *, include_url: bool = False) -> dict[str, Any]:
    out = {
        "job_id": job.get("job_id") or "",
        "thread_id": job.get("thread_id") or "",
        "host": job.get("host") or "",
        "status": job.get("status") or "",
        "created_at": job.get("created_at") or "",
        "claimed_at": job.get("claimed_at") or "",
        "terminal_at": job.get("terminal_at") or "",
        "error": job.get("error") or "",
        "punycode": bool(job.get("punycode")),
    }
    if include_url:
        out["url"] = job.get("url") or ""
    if job.get("result") is not None:
        out["result"] = copy.deepcopy(job.get("result"))
    return out


def pending_job_for_extension() -> dict[str, Any] | None:
    """Return the single pending job for the extension, if enabled."""
    with _LOCK:
        if not _AGENT_CAPTURE_ENABLED:
            raise BrowserCaptureError("capture_disabled", "Agent browser capture is disabled.", status_code=503)
        _expire_and_prune_locked()
        for job in _JOBS.values():
            if job.get("status") == STATUS_PENDING:
                return _public_job(job, include_url=True)
    return None


def claim_job(job_id: str) -> dict[str, Any]:
    """Atomically move a job from pending to claimed."""
    jid = str(job_id or "").strip()
    with _CONDITION:
        if not _AGENT_CAPTURE_ENABLED:
            raise BrowserCaptureError("capture_disabled", "Agent browser capture is disabled.", status_code=503)
        _expire_and_prune_locked()
        job = _JOBS.get(jid)
        if not job:
            raise BrowserCaptureError("not_found", "Capture job was not found.", status_code=404)
        status = job.get("status")
        if status == STATUS_EXPIRED:
            raise BrowserCaptureError("expired", "Capture job expired.", status_code=410)
        if status != STATUS_PENDING:
            raise BrowserCaptureError("invalid_state", "Capture job cannot be claimed.", status_code=409)
        job["status"] = STATUS_CLAIMED
        job["claimed_at"] = _now_iso()
        job["claimed_monotonic"] = time.time()
        _append_audit(job, STATUS_CLAIMED, decision="claimed")
        _CONDITION.notify_all()
        return _public_job(job, include_url=True)


def mark_complete(job_id: str, capture: dict[str, Any]) -> dict[str, Any]:
    """Complete a claimed job with captured content."""
    jid = str(job_id or "").strip()
    with _CONDITION:
        _expire_and_prune_locked()
        job = _JOBS.get(jid)
        if not job:
            raise BrowserCaptureError("not_found", "Capture job was not found.", status_code=404)
        status = job.get("status")
        if status == STATUS_EXPIRED:
            raise BrowserCaptureError("expired", "Capture job expired.", status_code=410)
        if _terminal(str(status)):
            raise BrowserCaptureError("invalid_state", "Capture job is already complete.", status_code=409)
        if status != STATUS_CLAIMED:
            raise BrowserCaptureError("invalid_state", "Capture job has not been claimed.", status_code=409)
        stored = dict(capture or {})
        stored["job_id"] = jid
        stored["captured_at"] = _now_iso()
        if "text_chars" not in stored:
            stored["text_chars"] = len(str(stored.get("text") or ""))
        job["status"] = STATUS_COMPLETE
        job["terminal_at"] = _now_iso()
        job["terminal_monotonic"] = time.time()
        job["decision"] = "user_allow"
        job["error"] = ""
        job["result"] = copy.deepcopy(stored)
        _append_audit(job, STATUS_COMPLETE, decision="user_allow")
        _CONDITION.notify_all()
        return _public_job(job, include_url=False)


def mark_failed(job_id: str, reason: str = "failed", *, decision: str = "user_cancel") -> dict[str, Any]:
    """Mark a pending or claimed job failed/cancelled."""
    jid = str(job_id or "").strip()
    safe_reason = str(reason or "failed").strip()[:80]
    with _CONDITION:
        _expire_and_prune_locked()
        job = _JOBS.get(jid)
        if not job:
            raise BrowserCaptureError("not_found", "Capture job was not found.", status_code=404)
        if _terminal(str(job.get("status"))):
            raise BrowserCaptureError("invalid_state", "Capture job is already terminal.", status_code=409)
        job["status"] = STATUS_FAILED
        job["terminal_at"] = _now_iso()
        job["terminal_monotonic"] = time.time()
        job["decision"] = str(decision or "user_cancel")[:40]
        job["error"] = safe_reason
        _append_audit(job, STATUS_FAILED, decision=job["decision"], error=safe_reason)
        _CONDITION.notify_all()
        return _public_job(job, include_url=False)


def get_job(job_id: str, *, thread_id: str | None = None, clear: bool = False) -> dict[str, Any] | None:
    """Return a job if it belongs to ``thread_id`` (when provided)."""
    jid = str(job_id or "").strip()
    with _CONDITION:
        _expire_and_prune_locked()
        job = _JOBS.get(jid)
        if not job:
            return None
        if thread_id is not None and str(job.get("thread_id") or "") != str(thread_id or ""):
            return None
        out = _public_job(job, include_url=False)
        if clear and _terminal(str(job.get("status"))):
            _JOBS.pop(jid, None)
        return out


def wait_for_job(job_id: str, *, thread_id: str | None = None, timeout: float = 120.0) -> dict[str, Any] | None:
    """Wait for a job to reach a terminal state, bounded by ``timeout`` seconds."""
    deadline = time.monotonic() + max(0.0, min(float(timeout or 0), 120.0))
    with _CONDITION:
        while True:
            _expire_and_prune_locked()
            job = _JOBS.get(str(job_id or "").strip())
            if job is None:
                return None
            if thread_id is not None and str(job.get("thread_id") or "") != str(thread_id or ""):
                return None
            if _terminal(str(job.get("status"))):
                return _public_job(job, include_url=False)
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                return _public_job(job, include_url=False)
            _CONDITION.wait(timeout=min(0.25, remaining))


def set_capture(capture: dict[str, Any]) -> dict[str, Any]:
    """Store the latest browser capture and return a defensive copy."""
    global _LATEST_CAPTURE
    stored = dict(capture or {})
    stored.setdefault("job_id", None)
    stored["captured_at"] = _now_iso()
    with _CONDITION:
        _LATEST_CAPTURE = copy.deepcopy(stored)
        _CONDITION.notify_all()
        return copy.deepcopy(_LATEST_CAPTURE)


def get_capture(*, clear: bool = True) -> dict[str, Any] | None:
    """Return the latest capture, optionally clearing it after read."""
    global _LATEST_CAPTURE
    with _LOCK:
        if _LATEST_CAPTURE is None:
            return None
        out = copy.deepcopy(_LATEST_CAPTURE)
        if clear:
            _LATEST_CAPTURE = None
        return out


def clear_capture() -> None:
    """Clear manual capture and all in-process browser-capture jobs."""
    global _LATEST_CAPTURE
    with _CONDITION:
        _LATEST_CAPTURE = None
        _JOBS.clear()
        _CREATE_JOB_HITS.clear()
        _CONDITION.notify_all()
