from __future__ import annotations

import os
import json
import signal
import string
import time
from pathlib import Path
from typing import Mapping

import yaml

from cybervisor.adapters.base import AgentAdapter
from cybervisor.adapters.types import LaunchRequest
from cybervisor.config import StageConfig, StageContractConfig, contract_artifact_path, contract_prompt_path
from cybervisor.hooks import set_active_stage_contract
from cybervisor.logging import CybervisorLogger, HookEventListener
from cybervisor.signals import SignalHandler


class PipelineInterrupted(Exception):
    def __init__(self, stage_name: str, attempt: int, exit_code: int | None = None) -> None:
        self.stage_name = stage_name
        self.attempt = attempt
        self.exit_code = exit_code
        detail = f" during attempt {attempt}" if attempt > 0 else ""
        status = f" (exit_code={exit_code})" if exit_code is not None else ""
        super().__init__(f"Pipeline interrupted in stage '{stage_name}'{detail}{status}")


def _is_interrupt_exit_code(exit_code: int | None) -> bool:
    if exit_code is None:
        return False
    if exit_code == 130:
        return True
    return exit_code in (-signal.SIGINT, -signal.SIGTERM)


def _render_prompt(template: str, context: Mapping[str, str]) -> tuple[str, set[str]]:
    formatter = string.Formatter()
    missing_keys: set[str] = set()
    for _literal_text, field_name, _format_spec, _conversion in formatter.parse(template):
        if field_name and field_name not in context:
            missing_keys.add(field_name)

    safe_context: dict[str, str] = dict(context)
    for key in missing_keys:
        safe_context[key] = ""
    return template.format_map(safe_context), missing_keys


def _artifact_file_path(stage_name: str) -> Path:
    return (Path.cwd() / contract_artifact_path(stage_name)).resolve()


def _prompt_file_path(stage_name: str) -> Path:
    return (Path.cwd() / contract_prompt_path(stage_name)).resolve()


def _missing_contract_prompt_paths(stages: list[StageConfig]) -> list[tuple[str, Path]]:
    missing: list[tuple[str, Path]] = []
    for stage in stages:
        if stage.contract is None:
            continue
        prompt_path = _prompt_file_path(stage.name)
        if not prompt_path.exists():
            missing.append((stage.name, prompt_path))
    return missing


def _validate_stage_artifact(stage_name: str, contract: StageContractConfig) -> tuple[dict[str, object] | None, str | None]:
    artifact_path = _artifact_file_path(stage_name)
    try:
        artifact_path.relative_to(Path.cwd().resolve())
    except ValueError:
        return None, "artifact_path must stay within the working directory"

    if not artifact_path.exists():
        return None, f"required artifact missing at {artifact_path}"

    try:
        payload = yaml.safe_load(artifact_path.read_text(encoding="utf-8"))
    except yaml.YAMLError as exc:
        return None, f"artifact at {artifact_path} is not valid YAML ({exc})"

    if not isinstance(payload, dict):
        return None, f"artifact at {artifact_path} must contain a YAML mapping"

    status = payload.get("status")
    if not isinstance(status, str) or status not in contract.allowed_statuses:
        return None, f"artifact status must be one of {', '.join(contract.allowed_statuses)}"

    summary = payload.get("summary")
    if not isinstance(summary, str) or not summary.strip():
        return None, "artifact summary must be a non-empty string"

    details = payload.get("details")
    if not isinstance(details, dict):
        return None, "artifact details must be a YAML mapping"

    actions = payload.get("actions")
    if not isinstance(actions, dict):
        return None, "artifact actions must be a YAML mapping"

    return payload, None


def _stringify_context_value(value: object) -> str:
    if value is None:
        return ""
    if isinstance(value, str):
        return value
    return yaml.safe_dump(value, sort_keys=False).rstrip()


def _stage_context_key(stage_name: str) -> str:
    return stage_name.lower().replace(" ", "_")


def _stage_log_path(stage_name: str) -> str:
    return f".cybervisor/logs/stages/{stage_name}.jsonl"


def _format_contract_artifact_output(stage_name: str, artifact_payload: dict[str, object]) -> str:
    artifact_yaml = yaml.safe_dump(artifact_payload, sort_keys=False).rstrip()
    return (
        f"[{stage_name}] Contract artifact output:\n"
        f"{contract_artifact_path(stage_name)}\n"
        f"{artifact_yaml}"
    )


class PipelineRunner:
    def __init__(self, log_file: str = ".cybervisor/logs/cybervisor.log.jsonl") -> None:
        self.log_file = log_file

    def _pop_hook_block(
        self,
        hook_listener: HookEventListener | None,
        agent_tool: str | None,
        *,
        wait_for_pending: bool = False,
    ) -> dict[str, object] | None:
        if hook_listener is None:
            return None

        hook_block = hook_listener.pop_blocking_decision(agent_tool)
        if hook_block is not None or not wait_for_pending:
            return hook_block

        deadline = time.monotonic() + 0.25
        while time.monotonic() < deadline:
            time.sleep(0.01)
            hook_block = hook_listener.pop_blocking_decision(agent_tool)
            if hook_block is not None:
                return hook_block
        return None

    def _build_prompt_context(
        self,
        prompt: str,
        stage_name: str,
        stage_context: Mapping[str, str],
    ) -> dict[str, str]:
        context = {
            "feature_description": prompt,
            "stage_name": stage_name,
            "spec_content": "",
            "plan_content": "",
            "implementation_summary": "",
            "test_plan": "",
            "latest_contract_yaml": "",
            "latest_contract_status": "",
            "latest_contract_summary": "",
            "latest_contract_details_yaml": "",
            "latest_contract_actions_yaml": "",
        }
        context.update(stage_context)
        return context

    def _build_injection_appendix(
        self,
        artifact_payload: dict[str, object],
        inject_sections: list[str],
    ) -> str:
        if not inject_sections:
            return ""

        parts: list[str] = []
        if "summary" in inject_sections:
            summary = _stringify_context_value(artifact_payload.get("summary"))
            if summary:
                parts.append(f"Summary:\n{summary}")
        if "details" in inject_sections:
            details_yaml = _stringify_context_value(artifact_payload.get("details"))
            if details_yaml:
                parts.append(f"Details:\n{details_yaml}")
        if "actions" in inject_sections:
            actions_yaml = _stringify_context_value(artifact_payload.get("actions"))
            if actions_yaml:
                parts.append(f"Required actions:\n{actions_yaml}")

        if not parts:
            return ""

        return "Latest routed contract context:\n\n" + "\n\n".join(parts)

    def _update_stage_context_from_artifact(
        self,
        stage: StageConfig,
        artifact_payload: dict[str, object],
        stage_context: dict[str, str],
        inject_sections: list[str],
    ) -> None:
        stage_key = _stage_context_key(stage.name)
        artifact_yaml = _stringify_context_value(artifact_payload)
        summary = _stringify_context_value(artifact_payload.get("summary"))
        details_yaml = _stringify_context_value(artifact_payload.get("details"))
        actions_yaml = _stringify_context_value(artifact_payload.get("actions"))
        status = str(artifact_payload.get("status", ""))

        stage_context["latest_contract_yaml"] = artifact_yaml
        stage_context["latest_contract_status"] = status
        stage_context["latest_contract_summary"] = summary if "summary" in inject_sections else ""
        stage_context["latest_contract_details_yaml"] = details_yaml if "details" in inject_sections else ""
        stage_context["latest_contract_actions_yaml"] = actions_yaml if "actions" in inject_sections else ""
        stage_context[f"{stage_key}_contract_yaml"] = artifact_yaml
        stage_context[f"{stage_key}_contract_status"] = status
        stage_context[f"{stage_key}_contract_summary"] = summary
        stage_context[f"{stage_key}_contract_details_yaml"] = details_yaml
        stage_context[f"{stage_key}_contract_actions_yaml"] = actions_yaml

    def _next_stage_index(
        self,
        stages: list[StageConfig],
        stage_lookup: Mapping[str, int],
        current_index: int,
        artifact_payload: dict[str, object] | None,
    ) -> tuple[int, str | None, str, list[str]]:
        stage = stages[current_index]
        status: str | None = None
        decision_source = "linear_fallback"

        if artifact_payload is not None:
            status = str(artifact_payload["status"])
            if stage.contract is not None and status in stage.contract.routes:
                route = stage.contract.routes[status]
                return stage_lookup[route.next_stage], status, "contract_route", list(route.injections)

        if stage.next_stage is not None:
            return stage_lookup[stage.next_stage], status, "success_route", []

        return current_index + 1, status, decision_source, []

    def run(
        self,
        prompt: str,
        stages: list[StageConfig],
        agent: AgentAdapter,
        logger: CybervisorLogger,
        signal_handler: SignalHandler | None = None,
        hook_listener: HookEventListener | None = None,
        start_stage_name: str | None = None,
        end_stage_name: str | None = None,
        end_before_stage_name: str | None = None,
    ) -> bool:
        stage_lookup = {stage.name: index for index, stage in enumerate(stages)}

        if end_stage_name is not None and end_stage_name not in stage_lookup:
            raise ValueError(f"Unknown end_stage_name '{end_stage_name}'. Valid stages are: {', '.join(stage_lookup.keys())}")
        if end_before_stage_name is not None and end_before_stage_name not in stage_lookup:
            raise ValueError(
                f"Unknown end_before_stage_name '{end_before_stage_name}'. Valid stages are: {', '.join(stage_lookup.keys())}"
            )
        if end_stage_name is not None and end_before_stage_name is not None:
            raise ValueError("--end-stage and --end-before cannot be used together.")

        missing_prompt_paths = _missing_contract_prompt_paths(stages)
        if missing_prompt_paths:
            details = "; ".join(f"{stage_name}: {path}" for stage_name, path in missing_prompt_paths)
            logger.log_to_stderr(f"[pipeline] Missing contract prompt files: {details}", "error")
            logger.log_to_json(
                logger.make_entry("Pipeline", "Failed", "Missing contract prompt files", error=details),
                self.log_file,
            )
            return False

        retry_counts = [0 for _ in stages]
        current_index = 0

        if start_stage_name is not None:
            current_index = stage_lookup[start_stage_name]

        if end_stage_name is not None:
            end_index = stage_lookup[end_stage_name]
            # Since linear pipelines can have loops, this is a coarse check for default sequences.
            if current_index > end_index:
                raise ValueError(f"Invalid sequence: start stage '{start_stage_name}' occurs after end stage '{end_stage_name}', so pipeline cannot realistically reach the end stage.")
            logger.log_to_json(
                {
                    "timestamp": logger.make_entry("Startup", "TargetedStart", "").timestamp,
                    "stage_name": "Startup",
                    "status": "TargetedStart",
                    "message": "Starting pipeline at requested stage",
                    "requested_stage_name": start_stage_name,
                    "effective_stage_name": stages[current_index].name,
                    "total_stage_count": len(stages),
                },
                self.log_file,
            )
        elif end_before_stage_name is not None:
            end_before_index = stage_lookup[end_before_stage_name]
            if current_index > end_before_index:
                raise ValueError(
                    f"Invalid sequence: start stage '{start_stage_name}' occurs after exclusive end stage '{end_before_stage_name}', so pipeline cannot stop before that stage."
                )
            if current_index == end_before_index:
                logger.log_to_json(
                    {
                        "timestamp": logger.make_entry("Startup", "TargetedEndBefore", "").timestamp,
                        "stage_name": "Startup",
                        "status": "TargetedEndBefore",
                        "message": "Skipping pipeline execution because start stage matches exclusive end boundary",
                        "requested_stage_name": end_before_stage_name,
                        "effective_stage_name": stages[current_index].name,
                        "total_stage_count": len(stages),
                    },
                    self.log_file,
                )
                return True

        stage_context: dict[str, str] = {}
        stage_appendices: dict[str, str] = {}

        while current_index < len(stages):
            if end_before_stage_name is not None and stages[current_index].name == end_before_stage_name:
                break
            stage = stages[current_index]
            attempt = retry_counts[current_index] + 1
            if signal_handler is not None and signal_handler.interrupted:
                raise PipelineInterrupted(stage.name, attempt - 1)
            if attempt > stage.max_retries:
                logger.log_to_stderr(f"[{stage.name}] Exhausted retries", "error")
                logger.log_to_json(
                    logger.make_entry(stage.name, "Failed", "Retries exhausted"),
                    self.log_file,
                )
                return False

            logger.log_to_stderr(f"[{stage.name}] Running attempt {attempt}")
            logger.log_to_json(
                logger.make_entry(stage.name, "Running", f"attempt={attempt}"),
                self.log_file,
            )

            context = self._build_prompt_context(prompt, stage.name, stage_context)
            stage_prompt, missing_keys = _render_prompt(stage.prompt_template, context)
            appended_context = stage_appendices.get(stage.name, "")
            if appended_context:
                stage_prompt = f"{stage_prompt.rstrip()}\n\n{appended_context}"
            if missing_keys:
                missing = ", ".join(sorted(missing_keys))
                logger.log_to_stderr(
                    f"[{stage.name}] Missing template placeholders defaulted to empty: {missing}",
                    "warning",
                )
            stage_output_log = _stage_log_path(stage.name)
            passed = False
            failure_error: str | None = None
            artifact_payload: dict[str, object] | None = None
            exit_code: int | None = None
            prompt_path = _prompt_file_path(stage.name) if stage.contract is not None else None
            if prompt_path is not None and not prompt_path.exists():
                logger.log_to_stderr(f"[{stage.name}] Missing contract prompt file: {prompt_path}", "error")
                logger.log_to_json(
                    logger.make_entry(stage.name, "Failed", "Missing contract prompt file"),
                    self.log_file,
                )
                return False
            try:
                if hook_listener is not None:
                    hook_listener.clear_pending_blocks(agent.descriptor.name)
                if stage.contract is not None:
                    artifact_path = _artifact_file_path(stage.name)
                    artifact_path.unlink(missing_ok=True)
                if agent.descriptor.capabilities.supports_hooks:
                    set_active_stage_contract(
                        agent.descriptor.name,
                        stage.name,
                        stage_prompt,
                        stage.contract,
                        attempt=attempt,
                        max_retries=stage.max_retries,
                    )
                running_process = agent.start(
                    LaunchRequest(
                        prompt=stage_prompt,
                        stage_name=stage.name,
                        log_file=stage_output_log,
                        logger=logger,
                    )
                )
                if signal_handler is not None:
                    signal_handler.set_pid(running_process.pid, running_process.process_group_id)
                exit_code, output = running_process.wait()
                if signal_handler is not None and signal_handler.interrupted:
                    raise PipelineInterrupted(stage.name, attempt, exit_code)
                if _is_interrupt_exit_code(exit_code):
                    raise PipelineInterrupted(stage.name, attempt, exit_code)
            except PipelineInterrupted:
                logger.log_to_stderr(f"[{stage.name}] Interrupted during attempt {attempt}", "warning")
                logger.log_to_json(
                    logger.make_entry(
                        stage.name,
                        "Interrupted",
                        "Stage interrupted by user",
                        agent_tool=agent.descriptor.name,
                        attempt=attempt,
                        exit_status=exit_code,
                        log_path=stage_output_log,
                    ),
                    self.log_file,
                )
                raise
            except Exception as exc:
                if signal_handler is not None and signal_handler.interrupted:
                    raise PipelineInterrupted(stage.name, attempt, exit_code) from exc
                failure_error = f"Stage execution failed: {exc}"
            else:
                passed = exit_code == 0
                hook_block = self._pop_hook_block(
                    hook_listener,
                    agent.descriptor.name,
                    wait_for_pending=passed,
                )
                if hook_block is not None:
                    passed = False
                    reason = hook_block.get("reason")
                    details = "Hook blocked agent completion"
                    if isinstance(reason, str) and reason:
                        details = f"{details}: {reason}"
                    failure_error = details
                elif stage.contract is not None and passed:
                    artifact_payload, artifact_error = _validate_stage_artifact(stage.name, stage.contract)
                    if artifact_error is not None:
                        passed = False
                        failure_error = f"Stage artifact validation failed: {artifact_error}"
                    elif artifact_payload is not None:
                        logger.log_to_stderr(_format_contract_artifact_output(stage.name, artifact_payload))
                        logger.log_to_json(
                            {
                                "timestamp": logger.make_entry(stage.name, "Execution", "").timestamp,
                                "stage_name": stage.name,
                                "status": "ArtifactValidated",
                                "message": "Validated stage artifact",
                                "artifact_path": str(contract_artifact_path(stage.name)),
                                "artifact_status": artifact_payload.get("status"),
                                "artifact_payload": artifact_payload,
                            },
                            self.log_file,
                        )
                elif not passed and failure_error is None and exit_code != 0:
                    failure_error = f"Agent exited with code {exit_code}"
            finally:
                if signal_handler is not None:
                    signal_handler.set_pid(None)
                if agent.descriptor.capabilities.supports_hooks:
                    set_active_stage_contract(agent.descriptor.name, None, None, None)

            if signal_handler is not None and signal_handler.interrupted:
                raise PipelineInterrupted(stage.name, attempt, exit_code)

            if not passed:
                retry_counts[current_index] += 1
                logger.log_to_stderr(f"[{stage.name}] Failed on attempt {attempt}", "warning")
                logger.log_to_json(
                    logger.make_entry(
                        stage.name,
                        "Failed",
                        "Stage failed, retrying",
                        error=failure_error,
                        agent_tool=agent.descriptor.name,
                        attempt=attempt,
                        exit_status=exit_code,
                        log_path=stage_output_log,
                    ),
                    self.log_file,
                )
                logger.log_to_json(
                    logger.make_entry(
                        stage.name,
                        "Execution",
                        "Stage attempt finished",
                        error=failure_error,
                        agent_tool=agent.descriptor.name,
                        attempt=attempt,
                        exit_status=exit_code,
                        log_path=stage_output_log,
                    ),
                    self.log_file,
                )
                continue

            next_index, status, decision_source, inject_sections = self._next_stage_index(
                stages,
                stage_lookup,
                current_index,
                artifact_payload,
            )
            if artifact_payload is not None:
                self._update_stage_context_from_artifact(stage, artifact_payload, stage_context, inject_sections)
            stage_appendices.pop(stage.name, None)

            if decision_source == "contract_route" and artifact_payload is not None:
                stage_appendices[stages[next_index].name] = self._build_injection_appendix(
                    artifact_payload,
                    inject_sections,
                )
            elif next_index < len(stages):
                stage_appendices.pop(stages[next_index].name, None)

            retry_counts[current_index] = 0

            logger.log_to_stderr(f"[{stage.name}] Success")
            logger.log_to_json(
                logger.make_entry(stage.name, "Success", "Stage completed"),
                self.log_file,
            )
            logger.log_to_json(
                logger.make_entry(
                    stage.name,
                    "Execution",
                    "Stage attempt finished",
                    agent_tool=agent.descriptor.name,
                    attempt=attempt,
                    exit_status=exit_code,
                    log_path=stage_output_log,
                ),
                self.log_file,
            )
            if artifact_payload is not None and status is not None:
                logger.log_to_json(
                    {
                        "timestamp": logger.make_entry(stage.name, "Routing", "").timestamp,
                        "stage_name": stage.name,
                        "status": "Routing",
                        "message": "Resolved next stage from contract status",
                        "artifact_status": status,
                        "next_stage": stages[next_index].name if next_index < len(stages) else None,
                        "decision_source": decision_source,
                    },
                    self.log_file,
                )

            if end_stage_name == stage.name:
                logger.log_to_stderr(f"[{stage.name}] Reached requested end_stage, stopping pipeline.")
                logger.log_to_json(
                    logger.make_entry(stage.name, "Execution", "Reached requested end_stage, stopping pipeline."),
                    self.log_file,
                )
                break

            current_index = next_index

        return True
