"""
Orchestrator Engine — composes all aura_core engines into a pipeline.

Stages run sequentially, each producing typed output that feeds the next.
Supports partial runs, stage skipping, and error-halt semantics.
Zero imports from app/. Uses module registry for engine dispatch.
"""
import time
import logging
from typing import Optional
from pydantic import BaseModel, Field

from aura_core.module_contract import aura_module, ModuleResult
from aura_core.types import OrchestratorInput

logger = logging.getLogger(__name__)

PIPELINE_STAGES = [
    {"name": "blueprint", "module": "blueprint", "order": 0, "required": True, "description": "Generate site blueprint with section definitions and depth config"},
    {"name": "valuation", "module": "valuation", "order": 1, "required": False, "description": "Calculate domain value, revenue models, and traffic estimates"},
    {"name": "context", "module": "context", "order": 2, "required": False, "description": "Assemble structured context for AI prompts"},
    {"name": "theme", "module": "theme", "order": 3, "required": True, "description": "Generate visual theme from brand colors and mood"},
    {"name": "validator", "module": "validator", "order": 4, "required": False, "description": "Validate and auto-repair generated content"},
]


class OrchestratorConfig(BaseModel):
    domain: str = ""
    depth: str = "comprehensive"
    skip_stages: list[str] = Field(default_factory=list)
    brand_colors: dict = Field(default_factory=dict)
    mood: str = "professional"
    niche: str = ""
    niches: list[dict] = Field(default_factory=list)


class OrchestratorResult(BaseModel):
    domain: str
    success: bool
    stages_completed: list[str] = Field(default_factory=list)
    stages_failed: list[str] = Field(default_factory=list)
    package: dict = Field(default_factory=dict)
    total_elapsed_ms: float = 0.0


class StageResult(BaseModel):
    ok: bool = True
    stage_name: str = ""
    output: Optional[dict] = None
    elapsed_ms: float = 0.0
    errors: list[str] = Field(default_factory=list)


def get_pipeline_config(
    domain: str = "",
    depth: str = "comprehensive",
    skip_stages: Optional[list[str]] = None,
    **kwargs,
) -> OrchestratorConfig:
    return OrchestratorConfig(
        domain=domain,
        depth=depth,
        skip_stages=skip_stages or [],
        **kwargs,
    )


def _run_blueprint(config: dict) -> dict:
    from aura_core.blueprint.engine import get_default_blueprint, validate_blueprint, blueprint_to_prompt_spec
    depth = config.get("depth", "comprehensive")
    bp = get_default_blueprint(depth)
    validation = validate_blueprint(bp)
    prompt_spec, json_schema = blueprint_to_prompt_spec(bp)
    return {
        "blueprint": bp.model_dump(),
        "validation": validation.model_dump(),
        "prompt_spec": prompt_spec,
        "json_schema": json_schema,
    }


def _run_valuation(config: dict) -> dict:
    from aura_core.valuation.engine import valuate_domain
    domain = config.get("domain", "")
    niches = config.get("niches", [])
    if not domain or not niches:
        return {"skipped": True, "reason": "domain or niches not provided"}
    result = valuate_domain(domain, {"niches": niches})
    return result.model_dump()


def _run_context(config: dict) -> dict:
    from aura_core.context.engine import build_domain_context
    domain = config.get("domain", "")
    analysis = config.get("analysis")
    niche = config.get("niche_data")
    brand = config.get("brand")
    valuation = config.get("valuation")
    context_str = build_domain_context(
        domain=domain,
        analysis=analysis,
        niche=niche,
        brand=brand,
        valuation=valuation,
    )
    return {"domain": domain, "context_text": context_str}


def _run_theme(config: dict) -> dict:
    from aura_core.theme.engine import generate_theme
    result = generate_theme(
        primary=config.get("primary_color", config.get("primary", "#4F46E5")),
        secondary=config.get("secondary_color", config.get("secondary", "#7C3AED")),
        accent=config.get("accent_color", config.get("accent", "#06B6D4")),
        mood=config.get("mood", "professional"),
        niche=config.get("niche", ""),
    )
    return result.model_dump()


def _run_validator(config: dict) -> dict:
    from aura_core.validator.engine import validate_site_copy
    site_copy = config.get("site_copy", {})
    repaired_data, reporter = validate_site_copy(site_copy)
    return reporter.to_dict()


_STAGE_RUNNERS = {
    "blueprint": _run_blueprint,
    "valuation": _run_valuation,
    "context": _run_context,
    "theme": _run_theme,
    "validator": _run_validator,
}


def run_stage(stage_name: str, config: dict) -> StageResult:
    runner = _STAGE_RUNNERS.get(stage_name)
    if runner is None:
        return StageResult(
            ok=False,
            stage_name=stage_name,
            errors=[f"Unknown stage: {stage_name}"],
        )
    start = time.perf_counter()
    try:
        output = runner(config)
        elapsed = (time.perf_counter() - start) * 1000
        return StageResult(
            ok=True,
            stage_name=stage_name,
            output=output,
            elapsed_ms=elapsed,
        )
    except Exception as e:
        elapsed = (time.perf_counter() - start) * 1000
        logger.exception(f"[orchestrator] Stage '{stage_name}' failed")
        return StageResult(
            ok=False,
            stage_name=stage_name,
            elapsed_ms=elapsed,
            errors=[str(e)],
        )


def validate_stage_output(stage_name: str, output) -> tuple[bool, list[str]]:
    if output is None:
        return False, [f"{stage_name}: output is None"]
    if not isinstance(output, dict):
        return False, [f"{stage_name}: output is not a dict"]
    return True, []


def assemble_package(domain: str, stage_results: dict[str, StageResult]) -> dict:
    package = {
        "domain": domain,
        "stages_completed": [],
        "stages_failed": [],
        "total_elapsed_ms": 0.0,
    }
    total_ms = 0.0
    for name, result in stage_results.items():
        total_ms += result.elapsed_ms
        if result.ok:
            package["stages_completed"].append(name)
            package[name] = result.output
        else:
            package["stages_failed"].append(name)
    package["total_elapsed_ms"] = round(total_ms, 2)
    return package


@aura_module(
    name="orchestrator",
    version="1.0.0",
    description="End-to-end pipeline composing all aura_core engines: blueprint → valuation → context → theme → validator",
    input_model=OrchestratorInput,
)
def run_pipeline(config: dict) -> ModuleResult:
    orch_config = OrchestratorConfig(**config) if not isinstance(config, OrchestratorConfig) else config
    stage_results = {}
    stages_completed = []
    stages_failed = []

    for stage_def in PIPELINE_STAGES:
        name = stage_def["name"]
        if name in orch_config.skip_stages:
            continue

        stage_config = _build_stage_config(name, orch_config, stage_results)
        result = run_stage(name, stage_config)
        stage_results[name] = result

        if result.ok:
            stages_completed.append(name)
        else:
            stages_failed.append(name)
            if stage_def["required"]:
                break

    package = assemble_package(orch_config.domain, stage_results)

    return ModuleResult(
        ok=len(stages_failed) == 0,
        output=OrchestratorResult(
            domain=orch_config.domain,
            success=len(stages_failed) == 0,
            stages_completed=stages_completed,
            stages_failed=stages_failed,
            package=package,
            total_elapsed_ms=package["total_elapsed_ms"],
        ),
    )


def _build_stage_config(stage_name: str, orch_config: OrchestratorConfig, prior_results: dict) -> dict:
    if stage_name == "blueprint":
        return {"depth": orch_config.depth}
    elif stage_name == "valuation":
        return {"domain": orch_config.domain, "niches": orch_config.niches}
    elif stage_name == "context":
        return {"domain": orch_config.domain, "niche": orch_config.niche}
    elif stage_name == "theme":
        colors = orch_config.brand_colors
        return {
            "primary_color": colors.get("primary", "#4F46E5"),
            "secondary_color": colors.get("secondary", "#7C3AED"),
            "accent_color": colors.get("accent", "#06B6D4"),
            "mood": orch_config.mood,
            "niche": orch_config.niche,
        }
    elif stage_name == "validator":
        bp_output = prior_results.get("blueprint", StageResult()).output or {}
        return {"site_copy": {}, "brand": {}}
    return {}