"""Execute a worker provider and capture logs/artifacts."""

from __future__ import annotations

import json
import socket
import shlex
import subprocess
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass, field, replace
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Optional

from loguru import logger

from ..utils import _now_iso
from ..worker import WorkerCancelledError, _run_codex_worker
from .config import WorkerProviderSpec


@dataclass(frozen=True)
class WorkerRunResult:
    """Structured outcome and artifacts from one worker command invocation."""
    provider: str
    prompt_path: str
    stdout_path: str
    stderr_path: str
    start_time: str
    end_time: str
    runtime_seconds: int
    exit_code: int
    timed_out: bool
    no_heartbeat: bool
    response_text: str = ""
    human_blocking_issues: list[dict[str, str]] = field(default_factory=list)
    token_usage: dict[str, Any] = field(default_factory=dict)


def _build_codex_command(spec: WorkerProviderSpec) -> str:
    """Build a codex command string with optional model/reasoning flags."""
    base = str(spec.command or "codex exec").strip() or "codex exec"
    parts = shlex.split(base)
    if not parts:
        parts = ["codex", "exec"]

    # Avoid duplicating flags if user already encoded them in command.
    has_model_flag = "--model" in parts
    has_reasoning_flag = "--reasoning-effort" in parts

    # Ensure non-interactive execution mode when the command does not already
    # specify a sandbox/approval mode.
    has_sandbox_flag = any(f in parts for f in ("--full-auto", "--sandbox", "--danger-full-access", "--yolo", "--dangerously-bypass-approvals-and-sandbox"))
    has_approval_flag = any(f in parts for f in ("--ask-for-approval", "-a"))
    if not has_sandbox_flag and not has_approval_flag:
        if spec.execution_mode == "host_access":
            parts.insert(1, "--danger-full-access")
        else:
            parts.insert(1, "--full-auto")

    if spec.model and not has_model_flag:
        parts.extend(["--model", str(spec.model)])
    supports_reasoning = _codex_supports_reasoning_effort(parts[0])
    if not supports_reasoning and parts[0] != "codex":
        supports_reasoning = _codex_supports_reasoning_effort("codex")
    if spec.reasoning_effort and not has_reasoning_flag and supports_reasoning:
        parts.extend(["--reasoning-effort", str(spec.reasoning_effort)])
    return shlex.join(parts)


def _build_claude_command(spec: WorkerProviderSpec) -> str:
    """Build a claude command string with optional model/effort flags."""
    base = str(spec.command or "claude -p").strip() or "claude -p"
    parts = shlex.split(base)
    if not parts:
        parts = ["claude", "-p"]
    if "-p" not in parts and "--print" not in parts:
        parts.append("-p")
    if spec.execution_mode == "host_access" and "--dangerously-skip-permissions" not in parts:
        parts.append("--dangerously-skip-permissions")

    has_model_flag = "--model" in parts
    has_effort_flag = "--effort" in parts

    if spec.model and not has_model_flag:
        parts.extend(["--model", str(spec.model)])
    supports_effort = _claude_supports_effort(parts[0])
    if not supports_effort and parts[0] != "claude":
        supports_effort = _claude_supports_effort("claude")
    if spec.reasoning_effort and not has_effort_flag and supports_effort:
        parts.extend(["--effort", str(spec.reasoning_effort)])
    output_format = _extract_option_value(parts, "--output-format")
    if not output_format:
        # Claude can stay silent in plain text mode for long intervals. Stream JSON
        # with partial chunks gives the orchestrator a steady liveness signal.
        if "--verbose" not in parts:
            parts.append("--verbose")
        parts.extend(["--output-format", "stream-json"])
        if "--include-partial-messages" not in parts:
            parts.append("--include-partial-messages")
    elif output_format == "stream-json":
        if "--verbose" not in parts:
            parts.append("--verbose")
        if "--include-partial-messages" not in parts:
            parts.append("--include-partial-messages")
    return shlex.join(parts)


def _extract_option_value(parts: list[str], option: str) -> str | None:
    for idx, part in enumerate(parts):
        if part == option and idx + 1 < len(parts):
            return str(parts[idx + 1]).strip().lower()
        if part.startswith(f"{option}="):
            return str(part.split("=", 1)[1]).strip().lower()
    return None


def _read_text(path_text: str) -> str:
    if not path_text:
        return ""
    try:
        return Path(path_text).read_text(encoding="utf-8", errors="replace")
    except Exception:
        return ""


def _extract_claude_stream_json_text(stdout_text: str) -> str:
    if not stdout_text.strip():
        return ""
    latest_assistant = ""
    result_text = ""
    delta_parts: list[str] = []
    for raw_line in stdout_text.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except Exception:
            continue
        if not isinstance(obj, dict):
            continue
        typ = str(obj.get("type") or "")
        if typ == "assistant":
            message = obj.get("message")
            if isinstance(message, dict):
                content = message.get("content")
                if isinstance(content, list):
                    texts: list[str] = []
                    for item in content:
                        if isinstance(item, dict) and str(item.get("type") or "") == "text":
                            text = item.get("text")
                            if isinstance(text, str):
                                texts.append(text)
                    merged = "".join(texts).strip()
                    if merged:
                        latest_assistant = merged
        elif typ == "result":
            maybe = obj.get("result")
            if isinstance(maybe, str) and maybe.strip():
                result_text = maybe.strip()
        elif typ == "stream_event":
            event = obj.get("event")
            if isinstance(event, dict) and str(event.get("type") or "") == "content_block_delta":
                delta = event.get("delta")
                if isinstance(delta, dict) and str(delta.get("type") or "") == "text_delta":
                    text = delta.get("text")
                    if isinstance(text, str) and text:
                        delta_parts.append(text)
    if latest_assistant:
        return latest_assistant
    if result_text:
        return result_text
    if delta_parts:
        return "".join(delta_parts).strip()
    return ""


def _extract_claude_stream_json_usage(stdout_text: str) -> dict[str, Any]:
    """Extract token usage from Claude CLI stream-json output.

    Parses NDJSON lines looking for ``{"type": "result", ...}`` events that
    contain ``usage`` and ``cost_usd`` fields.

    Args:
        stdout_text: Raw stdout captured from a Claude CLI invocation using
            ``--output-format stream-json``.

    Returns:
        A dict with ``input_tokens``, ``output_tokens``, ``cost_usd``, and
        ``provider_type`` keys when data is found; empty dict otherwise.
    """
    if not stdout_text.strip():
        return {}
    for raw_line in reversed(stdout_text.splitlines()):
        line = raw_line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except Exception:
            continue
        if not isinstance(obj, dict):
            continue
        if str(obj.get("type") or "") != "result":
            continue
        usage: dict[str, Any] = {"provider_type": "claude"}
        raw_usage = obj.get("usage")
        if isinstance(raw_usage, dict):
            for key in ("input_tokens", "output_tokens", "cache_creation_input_tokens", "cache_read_input_tokens"):
                val = raw_usage.get(key)
                if isinstance(val, (int, float)):
                    usage[key] = int(val)
        cost = obj.get("cost_usd")
        if isinstance(cost, (int, float)):
            usage["cost_usd"] = float(cost)
        return usage
    return {}


@lru_cache(maxsize=16)
def _codex_supports_reasoning_effort(executable: str) -> bool:
    """Best-effort check whether the codex CLI exposes --reasoning-effort."""
    try:
        completed = subprocess.run(
            [executable, "--help"],
            capture_output=True,
            text=True,
            check=False,
            timeout=5,
        )
    except (OSError, subprocess.SubprocessError):
        return False
    help_text = f"{completed.stdout}\n{completed.stderr}".lower()
    return "--reasoning-effort" in help_text


@lru_cache(maxsize=16)
def _claude_supports_effort(executable: str) -> bool:
    """Best-effort check whether the claude CLI exposes --effort."""
    try:
        completed = subprocess.run(
            [executable, "--help"],
            capture_output=True,
            text=True,
            check=False,
            timeout=5,
        )
    except (OSError, subprocess.SubprocessError):
        return False
    help_text = f"{completed.stdout}\n{completed.stderr}".lower()
    return "--effort" in help_text


def _extract_human_blocking_issues(progress_path: Path) -> list[dict[str, str]]:
    """Best-effort parse of human escalation issues from progress.json."""
    if not progress_path.exists():
        return []
    try:
        raw = json.loads(progress_path.read_text(encoding="utf-8"))
    except Exception:
        return []
    if not isinstance(raw, dict):
        return []
    issues = raw.get("human_blocking_issues")
    if not isinstance(issues, list):
        return []

    normalized: list[dict[str, str]] = []
    for item in issues:
        if isinstance(item, str):
            summary = item.strip()
            if summary:
                normalized.append({"summary": summary})
            continue

        if not isinstance(item, dict):
            continue

        summary = str(item.get("summary") or item.get("issue") or "").strip()
        details = str(item.get("details") or item.get("rationale") or "").strip()
        if not summary and details:
            summary = details.splitlines()[0][:200].strip()
        if not summary:
            continue

        issue: dict[str, str] = {"summary": summary}
        if details:
            issue["details"] = details
        for key in ("category", "action", "blocking_on", "severity"):
            value = item.get(key)
            if value is None:
                continue
            text = str(value).strip()
            if text:
                issue[key] = text
        normalized.append(issue)

    return normalized[:20]


def _run_ollama_generate(
    *,
    endpoint: str,
    model: str,
    prompt: str,
    run_dir: Path,
    timeout_seconds: int,
    temperature: Optional[float] = None,
    num_ctx: Optional[int] = None,
    is_cancelled: Optional[Callable[[], bool]] = None,
) -> WorkerRunResult:
    run_dir.mkdir(parents=True, exist_ok=True)
    prompt_path = run_dir / "prompt.txt"
    stdout_path = run_dir / "stdout.log"
    stderr_path = run_dir / "stderr.log"
    prompt_path.write_text(prompt)

    payload: dict[str, object] = {
        "model": model,
        "prompt": prompt,
        "stream": True,
    }
    options: dict[str, object] = {}
    if temperature is not None:
        options["temperature"] = float(temperature)
    if num_ctx is not None:
        options["num_ctx"] = int(num_ctx)
    if options:
        payload["options"] = options

    url = endpoint.rstrip("/") + "/api/generate"
    start_iso = _now_iso()
    start = time.monotonic()
    timed_out = False
    response_text_parts: list[str] = []
    ollama_token_usage: dict[str, Any] = {"provider_type": "ollama", "cost_usd": 0.0}

    request = urllib.request.Request(
        url,
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    # Use a small socket timeout so we can enforce overall timeout_seconds.
    socket_timeout_seconds = 15

    try:
        with urllib.request.urlopen(request, timeout=socket_timeout_seconds) as resp, open(
            stdout_path, "w", encoding="utf-8"
        ) as out, open(stderr_path, "w", encoding="utf-8") as err:
            while True:
                if timeout_seconds > 0 and time.monotonic() - start > timeout_seconds:
                    timed_out = True
                    err.write(f"[runner] Ollama timed out after {timeout_seconds}s\n")
                    break

                if is_cancelled and is_cancelled():
                    raise WorkerCancelledError("Task cancelled by user")

                try:
                    line = resp.readline()
                except socket.timeout:
                    continue

                if not line:
                    break
                try:
                    obj = json.loads(line.decode("utf-8", errors="replace"))
                except json.JSONDecodeError:
                    # Best-effort: log raw line and continue.
                    chunk = line.decode("utf-8", errors="replace")
                    err.write(chunk)
                    err.flush()
                    continue

                chunk = str(obj.get("response") or "")
                if chunk:
                    response_text_parts.append(chunk)
                    out.write(chunk)
                    out.flush()

                if bool(obj.get("done")):
                    prompt_eval = obj.get("prompt_eval_count")
                    eval_count = obj.get("eval_count")
                    if isinstance(prompt_eval, (int, float)):
                        ollama_token_usage["input_tokens"] = int(prompt_eval)
                    if isinstance(eval_count, (int, float)):
                        ollama_token_usage["output_tokens"] = int(eval_count)
                    break
    except urllib.error.HTTPError as exc:
        stderr_path.write_text(f"[runner] Ollama HTTP error: {exc.code} {exc.reason}\n")
        timed_out = False
    except urllib.error.URLError as exc:
        stderr_path.write_text(f"[runner] Ollama URL error: {exc.reason}\n")
        timed_out = False
    except Exception as exc:
        stderr_path.write_text(f"[runner] Ollama error: {exc.__class__.__name__}: {exc}\n")
        timed_out = False

    end_iso = _now_iso()
    runtime_seconds = int(time.monotonic() - start)
    response_text = "".join(response_text_parts)

    exit_code = 0
    if timed_out:
        exit_code = 124
    else:
        # Treat any stderr output as an error.
        try:
            if stderr_path.exists() and stderr_path.read_text(encoding="utf-8").strip():
                exit_code = 1
        except Exception:
            exit_code = 1

    return WorkerRunResult(
        provider=f"ollama:{model}",
        prompt_path=str(prompt_path),
        stdout_path=str(stdout_path),
        stderr_path=str(stderr_path),
        start_time=start_iso,
        end_time=end_iso,
        runtime_seconds=runtime_seconds,
        exit_code=exit_code,
        timed_out=timed_out,
        no_heartbeat=False,
        response_text=response_text,
        token_usage=ollama_token_usage,
    )


def run_worker(
    *,
    spec: WorkerProviderSpec,
    prompt: str,
    project_dir: Path,
    run_dir: Path,
    timeout_seconds: int,
    heartbeat_seconds: int,
    heartbeat_grace_seconds: int,
    progress_path: Path,
    expected_run_id: Optional[str] = None,
    on_spawn: Optional[Callable[[int], None]] = None,
    is_cancelled: Optional[Callable[[], bool]] = None,
    env: Optional[dict[str, str]] = None,
    cancel_event: Optional["threading.Event"] = None,
) -> WorkerRunResult:
    """Run the selected provider and return a normalized run result.

    Args:
        spec (WorkerProviderSpec): Spec for this call.
        prompt (str): Prompt for this call.
        project_dir (Path): Project dir for this call.
        run_dir (Path): Run dir for this call.
        timeout_seconds (int): Timeout seconds for this call.
        heartbeat_seconds (int): Heartbeat seconds for this call.
        heartbeat_grace_seconds (int): Heartbeat grace seconds for this call.
        progress_path (Path): Progress path for this call.
        expected_run_id (Optional[str]): Identifier for the related expected run.
        on_spawn (Optional[Callable[[int], None]]): On spawn for this call.
        is_cancelled (Optional[Callable[[], bool]]): Is cancelled for this call.
        env (Optional[dict[str, str]]): Environment variables for the subprocess.
            ``None`` inherits the parent process environment.
        cancel_event (Optional[threading.Event]): Event that, when set, unblocks
            the worker poll loop immediately for fast cancellation.

    Returns:
        WorkerRunResult: Result produced by this call.
    """
    if spec.type in {"codex", "claude"}:
        provider_label = "Codex" if spec.type == "codex" else "Claude"
        logger.info("Starting {} worker provider='{}' (timeout={}s)", provider_label, spec.name, timeout_seconds)
        command = _build_codex_command(spec) if spec.type == "codex" else _build_claude_command(spec)
        run_result = _run_codex_worker(
            command=command,
            prompt=prompt,
            project_dir=project_dir,
            run_dir=run_dir,
            timeout_seconds=timeout_seconds,
            heartbeat_seconds=heartbeat_seconds,
            heartbeat_grace_seconds=heartbeat_grace_seconds,
            progress_path=progress_path,
            expected_run_id=expected_run_id,
            on_spawn=on_spawn,
            is_cancelled=is_cancelled,
            env=env,
            cancel_event=cancel_event,
        )
        stdout_text = _read_text(str(run_result.get("stdout_path") or ""))
        response_text = stdout_text
        if spec.type == "claude" and _extract_option_value(shlex.split(command), "--output-format") == "stream-json":
            parsed = _extract_claude_stream_json_text(stdout_text)
            if parsed:
                response_text = parsed
        human_blocking_issues = _extract_human_blocking_issues(progress_path)
        # Extract token usage based on provider type
        token_usage: dict[str, Any] = {}
        if spec.type == "claude" and _extract_option_value(shlex.split(command), "--output-format") == "stream-json":
            token_usage = _extract_claude_stream_json_usage(stdout_text)
        elif spec.type == "codex":
            token_usage = {"provider_type": "codex"}
        return WorkerRunResult(
            provider=spec.name,
            prompt_path=str(run_result.get("prompt_path") or ""),
            stdout_path=str(run_result.get("stdout_path") or ""),
            stderr_path=str(run_result.get("stderr_path") or ""),
            start_time=str(run_result.get("start_time") or _now_iso()),
            end_time=str(run_result.get("end_time") or _now_iso()),
            runtime_seconds=int(run_result.get("runtime_seconds") or 0),
            exit_code=int(run_result.get("exit_code") or 0),
            timed_out=bool(run_result.get("timed_out")),
            no_heartbeat=bool(run_result.get("no_heartbeat")),
            response_text=response_text,
            human_blocking_issues=human_blocking_issues,
            token_usage=token_usage,
        )

    if spec.type == "ollama":
        logger.info(
            "Starting Ollama worker provider='{}' model='{}' (timeout={}s)",
            spec.name,
            spec.model,
            timeout_seconds,
        )
        result = _run_ollama_generate(
            endpoint=str(spec.endpoint),
            model=str(spec.model),
            prompt=prompt,
            run_dir=run_dir,
            timeout_seconds=timeout_seconds,
            temperature=spec.temperature,
            num_ctx=spec.num_ctx,
            is_cancelled=is_cancelled,
        )
        human_blocking_issues = _extract_human_blocking_issues(progress_path)
        if human_blocking_issues:
            return replace(result, human_blocking_issues=human_blocking_issues)
        return result

    raise ValueError(f"Unsupported worker type '{spec.type}'")
