from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
import logging
import subprocess

from rich.console import Console
from rich.panel import Panel
from armorer.agent_sdk import Status
from armorer.installers.schemas import DeploymentPlan, ExecutionResult
from armorer.scanner import RepoScanner
from armorer.agent import BaseAgent
from armorer.agent import shell as shell_tool, file_read, file_write
from armorer.security.policy import PolicyGuard
from armorer.tools.knowledge import search_knowledge, get_knowledge_topic
from armorer.tools.search import google_search, is_available as is_search_available
from armorer.tools.verify import verify_deployment

editor = None

logger = logging.getLogger(__name__)
console = Console()


def _activity_box(title: str, body: str):
    console.print()
    console.print(
        Panel(
            body,
            title=f"[bold cyan]{title}[/bold cyan]",
            border_style="cyan",
            expand=True,
        )
    )


@dataclass
class NodeResult:
    result: Any


@dataclass
class GraphExecutionResult:
    status: Status
    results: Dict[str, NodeResult] = field(default_factory=dict)


class InstallerGraph:
    """Deterministic installer orchestration with OpenAI Agents calls."""

    def __init__(self, builder: "InstallerGraphBuilder"):
        self.builder = builder

    def __call__(self, task: str) -> GraphExecutionResult:
        del task

        out: Dict[str, NodeResult] = {}
        base_agent = self.builder.base_agent

        try:
            if self.builder._should_use_repo_plan_directly():
                _activity_box("Armorer Activity", f"Using repository compose plan for [bold]{self.builder.app_name}[/bold].")
                try:
                    plan = self.builder._fallback_plan_from_repo()
                except ValueError:
                    if getattr(self.builder.installer_instance, "_active_playbook", None):
                        setup_steps = self.builder.installer_instance._detect_setup_sequence()
                        if setup_steps:
                            _activity_box(
                                "Armorer Activity",
                                "Generating playbook-managed runtime files before compose detection.",
                            )
                            self.builder.installer_instance._execute_setup_sequence(
                                setup_steps,
                                context="repo-plan-bootstrap",
                            )
                            plan = self.builder._fallback_plan_from_repo()
                        else:
                            raise
                    else:
                        raise
                out["repo_mapper"] = NodeResult(result=type("RepoMapResult", (), {"text": "deterministic-repo-plan"})())
                out["compose_planner"] = NodeResult(result=type("ArchitectResult", (), {"structured_output": plan})())
                out["hardening_reviewer"] = NodeResult(result=type("HardeningResult", (), {"structured_output": plan})())
                plan_dict = {
                    "docker-compose": plan.docker_compose,
                    "intent_checklist": plan.intent_checklist,
                    "run_commands": [],
                    "use_existing": True,
                }
                self.builder.installer_instance._prepare_files(plan_dict)
                setup_steps = self.builder.installer_instance._detect_setup_sequence()
                if setup_steps:
                    self.builder.installer_instance._execute_setup_sequence(setup_steps)
                self.builder.installer_instance._maybe_configure_env()
                deterministic_cmds = self.builder.installer_instance._verify_run_commands(
                    self.builder.installer_instance._playbook_run_commands()
                )
                self.builder.installer_instance._register(is_git=False, run_commands=deterministic_cmds)
                from armorer import docker_runner
                docker_runner.compose_down(
                    cwd=self.builder.installer_instance._docker_compose_cwd,
                    remove_orphans=True,
                )
                self.builder.installer_instance._remove_container_name_conflicts()
                console.print(f"  [dim]Using compose cwd: {self.builder.installer_instance._docker_compose_cwd}[/dim]")
                _activity_box("Armorer Activity", "Applying secure defaults and starting services.")
                self.builder.installer_instance._start_and_stabilize(max_attempts=3)
                final_exec = ExecutionResult(success=True, message="Provisioning complete from repository compose")
                out["stabilizer"] = NodeResult(result=type("StabilizerResult", (), {"structured_output": final_exec})())
                return GraphExecutionResult(status=Status.COMPLETED, results=out)

            # 1) Repo mapper
            _activity_box("Armorer Activity", "Analyzing repository structure and startup surfaces.")
            repo_map = base_agent.execute(
                self.builder.repo_mapper_agent,
                user_input=f"Map repository structure and startup surfaces for {self.builder.app_name}. Keep output concise.",
            )
            out["repo_mapper"] = NodeResult(result=type("RepoMapResult", (), {"text": str(repo_map)})())

            # 2) Compose planner
            _activity_box("Armorer Activity", "Planning Docker Compose deployment.")
            plan = None
            planner_prompt = (
                f"Perform deep research and planning for installation of {self.builder.app_name}. "
                f"Repository map:\n{repo_map}\n"
                "Return only the DeploymentPlan schema."
            )
            for _ in range(3):
                try:
                    candidate = base_agent.execute(
                        self.builder.compose_planner_agent,
                        user_input=planner_prompt,
                        structured_output_model=DeploymentPlan,
                    )
                except Exception:
                    candidate = None
                if isinstance(candidate, dict):
                    candidate = DeploymentPlan(**candidate)
                if isinstance(candidate, DeploymentPlan) and self.builder._is_valid_plan(candidate):
                    plan = candidate
                    break

            if plan is None:
                plan = self.builder._fallback_plan_from_repo()
            if not isinstance(plan, DeploymentPlan):
                if isinstance(plan, dict):
                    plan = DeploymentPlan(**plan)
                else:
                    raise ValueError(f"Architect did not return DeploymentPlan: {plan}")

            out["compose_planner"] = NodeResult(
                result=type("ArchitectResult", (), {"structured_output": plan})()
            )

            # 3) Hardening reviewer
            _activity_box("Armorer Activity", "Reviewing and hardening deployment configuration.")
            try:
                reviewed_plan = base_agent.execute(
                    self.builder.hardening_reviewer_agent,
                    user_input=(
                        "Review this DeploymentPlan for security and output a corrected DeploymentPlan.\n"
                        f"docker_compose:\n{plan.docker_compose}\n"
                        f"intent_checklist:{plan.intent_checklist}\n"
                    ),
                    structured_output_model=DeploymentPlan,
                )
            except Exception:
                reviewed_plan = plan
            if isinstance(reviewed_plan, dict):
                reviewed_plan = DeploymentPlan(**reviewed_plan)
            if isinstance(reviewed_plan, DeploymentPlan):
                plan = reviewed_plan
            out["hardening_reviewer"] = NodeResult(
                result=type("HardeningResult", (), {"structured_output": plan})()
            )

            # 4) Provision (deterministic)
            _activity_box("Armorer Activity", "Preparing files, env, and provisioning services.")
            plan_dict = {
                "docker-compose": plan.docker_compose,
                "intent_checklist": plan.intent_checklist,
                "run_commands": [],
                "use_existing": bool(self.builder.installer_instance._compose_dir),
            }
            self.builder.installer_instance._prepare_files(plan_dict)
            setup_steps = self.builder.installer_instance._detect_setup_sequence()
            if setup_steps:
                self.builder.installer_instance._execute_setup_sequence(setup_steps)
            self.builder.installer_instance._maybe_configure_env()
            deterministic_cmds = self.builder.installer_instance._verify_run_commands(
                self.builder.installer_instance._playbook_run_commands()
            )
            self.builder.installer_instance._register(is_git=False, run_commands=deterministic_cmds)
            from armorer import docker_runner
            docker_runner.compose_down(
                cwd=self.builder.installer_instance._docker_compose_cwd,
                remove_orphans=True,
            )
            self.builder.installer_instance._remove_container_name_conflicts()
            console.print(f"  [dim]Using compose cwd: {self.builder.installer_instance._docker_compose_cwd}[/dim]")
            self.builder.installer_instance._start_and_stabilize(max_attempts=3)

            provision_result = type(
                "ProvisionResult",
                (),
                {
                    "structured_output": ExecutionResult(success=True, message="Provisioning complete"),
                    "state": {"intent_checklist": plan.intent_checklist},
                },
            )()
            out["provisioner"] = NodeResult(result=provision_result)

            # 5) Stabilize loop
            _activity_box("Armorer Activity", "Validating runtime health and repairing if needed.")
            checklist = "\n".join(f"- {x}" for x in plan.intent_checklist)
            checklist_ctx = f"\n<intent_checklist>\n{checklist}\n</intent_checklist>\n" if checklist else ""

            final_exec: Optional[ExecutionResult] = None
            for _ in range(3):
                triage_res = base_agent.execute(
                    self.builder.triage_agent,
                    user_input=f"Verify deployment and completion of these goals:\n{checklist_ctx}",
                    structured_output_model=ExecutionResult,
                )
                if isinstance(triage_res, dict):
                    triage_res = ExecutionResult(**triage_res)
                if triage_res.success:
                    final_exec = triage_res
                    break

                diagnosis = base_agent.execute(
                    self.builder.log_agent,
                    user_input=f"Diagnose why system/checklist is failing:\n{checklist_ctx}",
                )
                base_agent.execute(
                    self.builder.repair_agent,
                    user_input=(
                        "Execute checklist and apply repairs based on diagnosis.\n"
                        f"Checklist:\n{checklist_ctx}\nDiagnosis:\n{diagnosis}"
                    ),
                )

            if final_exec is None:
                final_exec = ExecutionResult(success=False, message="Stabilization failed after retries")

            out["stabilizer"] = NodeResult(result=type("StabilizerResult", (), {"structured_output": final_exec})())
            return GraphExecutionResult(status=Status.COMPLETED if final_exec.success else Status.FAILED, results=out)

        except Exception as e:
            logger.error("Installer graph failed: %s", e)
            if "stabilizer" not in out:
                out["stabilizer"] = NodeResult(
                    result=type("StabilizerResult", (), {"structured_output": ExecutionResult(success=False, message=str(e))})()
                )
            return GraphExecutionResult(status=Status.FAILED, results=out)


class InstallerGraphBuilder:
    """Factory that creates a deterministic installer graph."""

    def __init__(
        self,
        installer_instance: Any,
        app_name: str,
        system_prompt: str,
        repo_docs: str = "",
        proactive_knowledge: str = "",
        status_obj: Any = None,
        provider: str = None,
        model_name: str = None,
        is_platform: bool = False,
        policy_guard: Optional[PolicyGuard] = None,
    ):
        self.installer_instance = installer_instance
        self.app_name = app_name
        self.system_prompt_helper = system_prompt
        self.repo_docs = repo_docs
        self.proactive_knowledge = proactive_knowledge
        self.status_obj = status_obj
        self.provider = provider
        self.model_name = model_name
        self.is_platform = is_platform
        self.policy_guard = policy_guard or PolicyGuard.default()

        self.base_agent: Optional[BaseAgent] = None
        self.repo_mapper_agent = None
        self.compose_planner_agent = None
        self.hardening_reviewer_agent = None
        self.triage_agent = None
        self.log_agent = None
        self.repair_agent = None

    def build(self, session_manager: Any = None) -> InstallerGraph:
        del session_manager

        base_agent = BaseAgent(self.provider, self.model_name)
        base_agent.set_policy_guard(self.policy_guard)
        self.base_agent = base_agent

        agent_protocol = (
            f"You are Armorer, an autonomous DevOps engineer. Project: {self.app_name}.\n"
            "OPERATIONAL RULES:\n"
            "1. NO GUESSING: use file_read/shell before concluding.\n"
            "2. OUTPUT STRICTNESS: follow schema exactly when requested.\n"
        )

        pre_scan_msg = ""
        try:
            repo_path = self.installer_instance.install_dir
            scanner = RepoScanner(repo_path)
            unique_files = scanner.scan_for_config_files()
            if unique_files:
                pre_scan_msg += "\n<pre_scanned_files>\n" + "\n".join([f"  - {f}" for f in unique_files]) + "\n</pre_scanned_files>\n"
        except Exception:
            unique_files = []

        repo_docs_section = f"\n<repository_documentation>\n{self.repo_docs}\n</repository_documentation>\n" if self.repo_docs else ""

        discovery_tools = [file_read, shell_tool, search_knowledge]
        if is_search_available():
            discovery_tools.append(google_search)

        self.repo_mapper_agent = base_agent.get_agent(
            agent_name="repo-mapper",
            system_prompt=(
                f"{agent_protocol}\n"
                "ROLE: Repo Mapper\n"
                "GOAL: map files, compose locations, env files, startup scripts.\n"
                "Keep output terse and factual."
            ),
            tools=discovery_tools,
            status_obj=self.status_obj,
            model=self.model_name,
        )

        planner_instructions = (
            "ROLE: Compose Planner\n"
            "GOAL: produce a minimal DeploymentPlan schema suitable for small models.\n"
            "Use intent_checklist (not run_commands). Return strict schema only."
        )

        self.compose_planner_agent = base_agent.get_agent(
            agent_name="compose-planner",
            system_prompt=f"{agent_protocol}\n{planner_instructions}\n{pre_scan_msg}\n{repo_docs_section}\n{self.system_prompt_helper}",
            tools=discovery_tools,
            status_obj=self.status_obj,
            model=self.model_name,
            structured_output_model=DeploymentPlan,
        )

        self.hardening_reviewer_agent = base_agent.get_agent(
            agent_name="hardening-reviewer",
            system_prompt=(
                f"{agent_protocol}\n"
                "ROLE: Hardening Reviewer\n"
                "GOAL: enforce secure docker defaults and return corrected DeploymentPlan.\n"
                "Do not add unnecessary services."
            ),
            tools=[file_read, search_knowledge, get_knowledge_topic],
            status_obj=self.status_obj,
            model=self.model_name,
            structured_output_model=DeploymentPlan,
        )

        common_tools = [shell_tool, file_read, file_write, search_knowledge, get_knowledge_topic]
        if is_search_available():
            common_tools.append(google_search)

        self.triage_agent = base_agent.get_agent(
            agent_name="triage",
            system_prompt=(
                f"{agent_protocol}\n"
                "ROLE: Triage Commander\n"
                "GOAL: verify system health and checklist completion.\n"
                "Use verify_deployment first."
            ),
            tools=[verify_deployment] + common_tools,
            status_obj=self.status_obj,
            model=self.model_name,
            structured_output_model=ExecutionResult,
        )

        self.log_agent = base_agent.get_agent(
            agent_name="log_analyst",
            system_prompt=(
                f"{agent_protocol}\n"
                "ROLE: Log Detective\n"
                "GOAL: diagnose deployment failures precisely."
            ),
            tools=common_tools,
            status_obj=self.status_obj,
            model=self.model_name,
        )

        sage_tools = common_tools + ([editor] if editor else [])
        self.repair_agent = base_agent.get_agent(
            agent_name="repair-agent",
            system_prompt=(
                f"{agent_protocol}\n"
                "ROLE: Repair Agent\n"
                "GOAL: execute checklist and apply bounded, deterministic repairs."
            ),
            tools=sage_tools,
            status_obj=self.status_obj,
            model=self.model_name,
        )

        return InstallerGraph(self)

    def _is_valid_plan(self, plan: DeploymentPlan) -> bool:
        try:
            import yaml

            data = yaml.safe_load(plan.docker_compose)
            if not isinstance(data, dict):
                return False
            services = data.get("services")
            if not isinstance(services, dict) or not services:
                return False
            return True
        except Exception:
            return False

    def _fallback_plan_from_repo(self) -> DeploymentPlan:
        """Deterministic fallback when small model produces invalid schema content."""
        compose_candidates = []
        seen = set()

        def add_candidate(p):
            try:
                rp = p.resolve()
            except Exception:
                rp = p
            key = str(rp)
            if key not in seen:
                seen.add(key)
                compose_candidates.append(p)

        if "autogpt" in self.app_name.lower():
            add_candidate(self.installer_instance.install_dir / "autogpt_platform" / "docker-compose.yml")
            add_candidate(self.installer_instance.install_dir / "autogpt_platform" / "docker-compose.platform.yml")
            add_candidate(self.installer_instance.install_dir / "classic" / "original_autogpt" / "docker-compose.yml")
            add_candidate(self.installer_instance.install_dir / "classic" / "docker-compose.yml")

        add_candidate(self.installer_instance.install_dir / "docker-compose.yml")
        add_candidate(self.installer_instance.install_dir / "docker-compose.yaml")
        add_candidate(self.installer_instance.install_dir / "autogpt_platform" / "docker-compose.yml")
        add_candidate(self.installer_instance.install_dir / "autogpt_platform" / "docker-compose.platform.yml")

        # Generic fallback: discover compose files recursively in repo.
        # Prefer shallower paths first to avoid selecting deeply nested examples.
        try:
            discovered = []
            for p in self.installer_instance.install_dir.rglob("docker-compose*.yml"):
                if p.is_file():
                    discovered.append(p)
            for p in self.installer_instance.install_dir.rglob("docker-compose*.yaml"):
                if p.is_file():
                    discovered.append(p)
            discovered.sort(key=lambda x: (len(x.relative_to(self.installer_instance.install_dir).parts), str(x)))
            for p in discovered:
                add_candidate(p)
        except Exception:
            pass

        for candidate in compose_candidates:
            if candidate.exists():
                raw = candidate.read_text()
                try:
                    content = self.installer_instance._normalize_compose_yaml(raw)
                except Exception:
                    content = raw
                plan = DeploymentPlan(
                    docker_compose=content,
                    intent_checklist=[
                        "Ensure required .env files are present",
                        "Bring up stack with docker compose",
                    ],
                )
                if self._is_valid_plan(plan) and not self._has_missing_extends(content, candidate.parent):
                    self.installer_instance._compose_dir = candidate.parent
                    return plan
        raise ValueError("Unable to create valid fallback DeploymentPlan from repository.")

    def _has_missing_extends(self, content: str, base_dir: Any) -> bool:
        try:
            import yaml

            data = yaml.safe_load(content) or {}
            services = data.get("services", {})
            for svc in services.values():
                ext = svc.get("extends") if isinstance(svc, dict) else None
                if isinstance(ext, dict) and ext.get("file"):
                    path = (base_dir / ext["file"]).resolve()
                    if not path.exists():
                        return True
            return False
        except Exception:
            return True

    def _should_use_repo_plan_directly(self) -> bool:
        # Deterministic fast path for repositories that already ship compose stacks.
        if getattr(self.installer_instance, "_active_playbook", None):
            return True
        markers = [
            self.installer_instance.install_dir / "docker-compose.yml",
            self.installer_instance.install_dir / "docker-compose.yaml",
            self.installer_instance.install_dir / "autogpt_platform" / "docker-compose.yml",
            self.installer_instance.install_dir / "classic" / "original_autogpt" / "docker-compose.yml",
        ]
        if any(p.exists() for p in markers):
            return True
        try:
            for _ in self.installer_instance.install_dir.rglob("docker-compose*.yml"):
                return True
            for _ in self.installer_instance.install_dir.rglob("docker-compose*.yaml"):
                return True
        except Exception:
            pass
        return False
