#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional, Tuple

import numpy as np

from forecast_diagnostics import diebold_mariano_squared_error, rolling_mean_forecast
from sprint_forecast import (
    BiasModelParameters,
    BlockedStateSpec,
    SprintPlan,
    SprintSuccessForecaster,
    WorkItemEstimate,
)
from sprint_forecast.forecaster import DemandMode
from zenhub_client import ZenhubClient, ZenhubError


def _load_config(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _as_2d_matrix(value: Any, *, n_rows: int, n_cols: int, name: str) -> np.ndarray:
    """Coerce a scalar, 1D, or 2D config field to a (n_rows, n_cols) float matrix."""
    if value is None:
        raise ValueError(f"{name} is required")
    arr = np.asarray(value, dtype=float)
    if arr.ndim == 0:
        return np.full((n_rows, n_cols), float(arr))
    if arr.ndim == 1:
        if arr.shape[0] != n_rows:
            raise ValueError(f"{name} must have length {n_rows} when 1D")
        return np.repeat(arr.reshape(n_rows, 1), n_cols, axis=1)
    if arr.ndim == 2:
        if arr.shape != (n_rows, n_cols):
            raise ValueError(f"{name} must have shape {(n_rows, n_cols)}")
        return arr
    raise ValueError(f"{name} must be scalar, 1D, or 2D")


def _build_plan(cfg: Dict[str, Any]) -> SprintPlan:
    devs = cfg.get("developers", [])
    if not devs:
        raise ValueError("Config must include non-empty 'developers' list")

    developer_ids = [str(d["name"]) for d in devs]
    n = len(developer_ids)

    T = int(cfg.get("sprint_days", cfg.get("days", 10)))
    if T <= 0:
        raise ValueError("sprint_days must be positive")

    h_per_day = np.asarray([float(d["h"]) for d in devs], dtype=float)
    q_per_day = np.asarray([float(d["q"]) for d in devs], dtype=float)
    mean_focus = np.asarray([float(d["mu_F"]) for d in devs], dtype=float)

    if np.any(mean_focus <= 0):
        raise ValueError("developers[*].mu_F must be positive")

    sigma_F = cfg.get("sigma_F")
    if sigma_F is None:
        sigma_per_dev = np.full(n, 0.25, dtype=float)
    else:
        sigma_per_dev = np.asarray(sigma_F, dtype=float).reshape(-1)
        if sigma_per_dev.size != n:
            raise ValueError("sigma_F must match number of developers")

    planned_hours = np.repeat(h_per_day.reshape(n, 1), T, axis=1)
    availability_prob = np.repeat(q_per_day.reshape(n, 1), T, axis=1)

    # Map desired mean of a lognormal to log-mu given sigma:
    # mean = exp(mu + sigma^2/2) -> mu = log(mean) - sigma^2/2.
    focus_log_sigma = np.repeat(sigma_per_dev.reshape(n, 1), T, axis=1)
    focus_log_mu = np.repeat(
        (np.log(mean_focus) - 0.5 * (sigma_per_dev**2)).reshape(n, 1), T, axis=1
    )

    u_cfg = cfg.get("dev_day_ai_intensity", 0.0)
    dev_day_ai_intensity = _as_2d_matrix(u_cfg, n_rows=n, n_cols=T, name="dev_day_ai_intensity")

    blocked_cfg = cfg.get("blocked")
    if blocked_cfg:
        g = np.asarray(blocked_cfg.get("multipliers"), dtype=float).reshape(-1)
        pi_raw = blocked_cfg.get("probabilities")
        pi = np.asarray(pi_raw, dtype=float)
        if pi.ndim == 1:
            if pi.shape[0] != g.shape[0]:
                raise ValueError("blocked.probabilities length must match multipliers")
            pi = np.repeat(pi.reshape(g.shape[0], 1), T, axis=1)
        blocked = BlockedStateSpec(multipliers=g, probabilities=pi)
    else:
        blocked = BlockedStateSpec(multipliers=np.array([1.0]), probabilities=np.ones((1, T)))

    return SprintPlan(
        developer_ids=developer_ids,
        planned_hours=planned_hours,
        availability_prob=availability_prob,
        dev_day_ai_intensity=dev_day_ai_intensity,
        focus_log_mu=focus_log_mu,
        focus_log_sigma=focus_log_sigma,
        blocked=blocked,
    )


def _maybe_get_workload_from_zenhub(cfg: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    workload_source = cfg.get("workload_source")
    if not workload_source:
        return None

    if workload_source.get("type") != "zenhub_sprint_points":
        return None

    token = os.getenv("ZENHUB_API_TOKEN") or workload_source.get("api_token")
    if not token:
        raise RuntimeError("ZENHUB_API_TOKEN env var (or workload_source.api_token) is required for Zenhub workload pull.")

    workspace_id = workload_source.get("workspace_id")
    if not workspace_id:
        raise ValueError("workload_source.workspace_id is required")

    sprint_id = workload_source.get("sprint_id")
    use_current_open = bool(workload_source.get("use_current_open_sprint", True))

    zh = ZenhubClient(api_token=token)

    if not sprint_id and use_current_open:
        sprint = zh.sprints.get_current_open_sprint(workspace_id)
        sprint_id = sprint.id

    if not sprint_id:
        raise ValueError("Provide workload_source.sprint_id, or set use_current_open_sprint=true with a valid workspace_id.")

    summary = zh.sprints.get_sprint_workload_points(sprint_id)

    return {
        "workload": summary.total_points,
        "workload_units": "story_points",
        "sprint": {
            "id": summary.sprint.id,
            "name": summary.sprint.name,
            "state": summary.sprint.state,
            "startAt": summary.sprint.start_at,
            "endAt": summary.sprint.end_at,
        },
        "issue_count": summary.issue_count,
        "missing_estimate_count": summary.missing_estimate_count,
        "closed_issue_count": summary.closed_issue_count,
        "closed_points": summary.closed_points,
    }


def _build_work_items(cfg: Dict[str, Any], *, fallback_owner_id: str) -> Tuple[list[WorkItemEstimate], Optional[Dict[str, Any]]]:
    """Build work items from config; fall back to a single aggregate item."""
    work_items = cfg.get("work_items")
    if isinstance(work_items, list) and work_items:
        items: list[WorkItemEstimate] = []
        for i, raw in enumerate(work_items):
            if not isinstance(raw, dict):
                raise ValueError("work_items entries must be objects")
            item_id = str(raw.get("id", raw.get("item_id", f"item-{i+1}")))
            owner = str(raw.get("owner_id", raw.get("engineer", raw.get("owner", fallback_owner_id))))
            o = float(raw["optimistic"])
            m = float(raw["most_likely"])
            p = float(raw["pessimistic"])
            ai = float(raw.get("ai_intensity", 0.0))
            items.append(WorkItemEstimate(item_id, owner, o, m, p, ai_intensity=ai))
        return items, {"source": "work_items", "n_items": len(items)}

    workload = cfg.get("workload")
    workload_meta = None
    if workload is None:
        workload_meta = _maybe_get_workload_from_zenhub(cfg)
        if workload_meta is not None:
            workload = workload_meta["workload"]

    # New: build per-issue work items from Zenhub sprint + PERT export.
    workload_source = cfg.get("workload_source")
    if workload is None and isinstance(workload_source, dict) and workload_source.get("type") == "zenhub_sprint_issues":
        token = os.getenv("ZENHUB_API_TOKEN") or workload_source.get("api_token")
        if not token:
            raise RuntimeError(
                "ZENHUB_API_TOKEN env var (or workload_source.api_token) is required for Zenhub sprint issue pull."
            )
        zh = ZenhubClient(api_token=token)
        from sprint_forecast.integrations.zenhub_sprint_issues import build_work_items_from_zenhub_sprint

        result = build_work_items_from_zenhub_sprint(
            zh=zh,
            source_cfg=workload_source,
            fallback_owner_id=fallback_owner_id,
            developer_ids=cfg.get("developers") and [str(d["name"]) for d in cfg.get("developers", [])],
        )

        meta = dict(result.meta)
        meta["project_items"] = [pi.to_json() for pi in result.project_items]
        return result.items, meta

    if workload is None:
        raise ValueError(
            "Provide either 'work_items' (preferred) or a scalar 'workload' / workload_source.type='zenhub_sprint_points'."
        )

    w = float(workload)
    if w <= 0:
        raise ValueError("workload must be positive")

    # Create a single aggregate item with a narrow triad.
    o = 0.9 * w
    m = w
    p = 1.1 * w
    items = [WorkItemEstimate("aggregate", fallback_owner_id, o, m, p, ai_intensity=0.0)]
    meta = {"source": "aggregate_workload", "workload": w}
    if workload_meta is not None:
        meta["workload_meta"] = workload_meta
    return items, meta


def _format_markdown_summary(payload: Dict[str, Any]) -> str:
    """
    A Zenhub-friendly status summary suitable for a "project status" issue.
    """
    p = payload["p_success"]

    def risk_badge(p_success: float) -> str:
        if p_success >= 0.85:
            return "🟢 High"
        if p_success >= 0.65:
            return "🟡 Medium"
        return "🔴 Low"

    lines = []
    lines.append(f"## Sprint delivery forecast ({payload['generated_at']})")
    lines.append("")
    lines.append(f"**P(success):** **{p:.1%}** ({risk_badge(p)})")
    lines.append("")
    lines.append("| Quantity | p10 | p50 | p90 | mean |")
    lines.append("| --- | --- | --- | --- | --- |")
    cap = payload["capacity_summary"]
    wl = payload["workload_summary"]
    lines.append(
        f"| Remaining capacity | {cap['p10']:.2f} | {cap['p50']:.2f} | {cap['p90']:.2f} | {cap['mean']:.2f} |"
    )
    lines.append(
        f"| Remaining workload | {wl['p10']:.2f} | {wl['p50']:.2f} | {wl['p90']:.2f} | {wl['mean']:.2f} |"
    )
    lines.append("")
    if payload.get("work_items_meta", {}).get("workload_meta", {}).get("sprint"):
        meta = payload["work_items_meta"]["workload_meta"]
        lines.append("### Workload source (Zenhub)")
        lines.append(f"- Sprint: **{meta.get('sprint',{}).get('name','(unknown)')}** (`{meta.get('sprint',{}).get('id','')}`)")
        lines.append(f"- Issues: {meta.get('issue_count',0)}, missing estimates: {meta.get('missing_estimate_count',0)}")
        lines.append(f"- Closed issues: {meta.get('closed_issue_count',0)}, closed points: {meta.get('closed_points',0.0):.2f}")
        lines.append("")

    if payload.get("ai_uplift"):
        u = payload["ai_uplift"]
        lines.append("### AI counterfactual")
        lines.append(f"- P(success | AI)   = {u['p_ai']:.1%}")
        lines.append(f"- P(success | noAI) = {u['p_no_ai']:.1%}")
        lines.append(f"- Delta p           = {u['delta_p_overall']:+.1%}")
        lines.append("")

    lines.append("### Notes")
    lines.append("- This script uses the exact-state simulation forecaster (see `examples/demo.py`).")
    lines.append("- Prefer `work_items` (three-point estimates) over a single scalar workload.")
    return "\n".join(lines) + "\n"


def _maybe_update_zenhub_status(cfg: Dict[str, Any], markdown: str) -> Optional[Dict[str, Any]]:
    status_cfg = cfg.get("zenhub_status")
    if not status_cfg or not status_cfg.get("enabled", False):
        return None

    token = os.getenv("ZENHUB_API_TOKEN") or status_cfg.get("api_token")
    if not token:
        raise RuntimeError("ZENHUB_API_TOKEN env var (or zenhub_status.api_token) is required to update Zenhub status issue.")

    workspace_id = status_cfg.get("workspace_id")
    repo_gh_id = status_cfg.get("status_issue_repository_gh_id")
    issue_number = status_cfg.get("status_issue_number")

    if not workspace_id or repo_gh_id is None or issue_number is None:
        raise ValueError("zenhub_status must include workspace_id, status_issue_repository_gh_id, status_issue_number")

    zh = ZenhubClient(api_token=token)
    issue = zh.issues.get_issue_by_info(repository_gh_id=int(repo_gh_id), issue_number=int(issue_number))
    zh.issues.update_issue(issue_id=issue.id, body=markdown)

    # Optional move to a pipeline (nice for dashboards)
    pipeline_name = status_cfg.get("pipeline_name")
    if pipeline_name:
        pipeline_id = zh.pipelines.find_pipeline_id(workspace_id, pipeline_name)
        zh.pipelines.move_issue(workspace_id=workspace_id, issue_id=issue.id, pipeline_id=pipeline_id, position=0)

    return {
        "status_issue_id": issue.id,
        "status_issue_number": issue.number,
        "status_issue_title": issue.title,
        "status_issue_html_url": issue.html_url,
        "pipeline_name": pipeline_name,
    }


def _maybe_compute_dm(cfg: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    dm_cfg = cfg.get("forecast_accuracy")
    if not dm_cfg or not dm_cfg.get("enabled", False):
        return None

    path = Path(dm_cfg.get("history_json", dm_cfg.get("history_path", "")))
    if not path.exists():
        raise FileNotFoundError(f"forecast_accuracy history file not found: {path}")

    window = int(dm_cfg.get("window_k", 5))

    # Expect JSON array of { "actual": ..., "sc_forecast": ... } entries
    series = json.loads(path.read_text(encoding="utf-8"))
    actual = np.asarray([float(x["actual"]) for x in series], dtype=float)
    sc = np.asarray([float(x["sc_forecast"]) for x in series], dtype=float)

    idx, rm_forecasts = rolling_mean_forecast(actual, window=window)
    actual_aligned = actual[idx]
    sc_aligned = sc[idx]

    dm = diebold_mariano_squared_error(actual_aligned, forecast_sc=sc_aligned, forecast_rm=rm_forecasts)

    return {
        "window_k": window,
        "aligned_T": int(dm.T),
        "dm_test": dm.to_json(),
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Sprint success forecast (exact-state, simulation-first).")
    ap.add_argument("--config", default="config.json", help="Path to config JSON.")
    ap.add_argument("--outdir", default="forecast_output", help="Output directory for forecast.json and forecast.md")
    ap.add_argument("--completed-days", type=int, default=None, help="Completed days t (overrides config if set).")
    ap.add_argument("--observed-cum-capacity", type=float, default=None, help="Observed cumulative capacity Ĉ_t (overrides config if set).")
    ap.add_argument("--demand-mode", choices=[m.value for m in DemandMode], default=None, help="Demand interpretation (overrides config if set).")
    args = ap.parse_args()

    cfg = _load_config(Path(args.config))

    units = cfg.get("units", "units")
    monte_carlo_samples = int(cfg.get("monte_carlo_samples", 100_000))
    seed = cfg.get("seed")

    plan = _build_plan(cfg)

    bias_cfg = cfg.get("bias") or {}
    bias = BiasModelParameters(
        mu_alpha=float(bias_cfg.get("mu_alpha", 0.0)),
        beta_ai=float(bias_cfg.get("beta_ai", 0.0)),
        alpha_by_engineer=dict(bias_cfg.get("alpha_by_engineer", {})),
        delta_b=float(bias_cfg.get("delta_b", 0.10)),
    )
    beta_ai_focus = float(cfg.get("beta_ai_focus", 0.0))
    forecaster = SprintSuccessForecaster(plan, bias=bias, beta_ai_focus=beta_ai_focus)

    completed_days = int(args.completed_days if args.completed_days is not None else cfg.get("completed_days", 0))
    observed_cum_capacity = float(
        args.observed_cum_capacity
        if args.observed_cum_capacity is not None
        else cfg.get("observed_cum_capacity", 0.0)
    )
    demand_mode = DemandMode(
        args.demand_mode if args.demand_mode is not None else cfg.get("demand_mode", DemandMode.REMAINING_ONLY.value)
    )

    items, work_items_meta = _build_work_items(cfg, fallback_owner_id=str(plan.developer_ids[0]))

    # Monte Carlo summaries.
    rng = np.random.default_rng(seed)
    workload_sampler = forecaster.workload_builder.build_sampler(items, bias)
    w = workload_sampler.sample_total(rng=rng, size=monte_carlo_samples)
    s = forecaster.capacity_sim.simulate_remaining_capacity(
        completed_days=completed_days, n_samples=monte_carlo_samples, rng=rng
    ).samples

    def _summ(x: np.ndarray) -> Dict[str, float]:
        qs = np.quantile(x, [0.10, 0.50, 0.90])
        return {
            "mean": float(np.mean(x)),
            "p10": float(qs[0]),
            "p50": float(qs[1]),
            "p90": float(qs[2]),
        }

    workload_summary = _summ(w)
    capacity_summary = _summ(s)

    p_success = forecaster.estimate_success_probability(
        completed_days=completed_days,
        observed_cum_capacity=observed_cum_capacity,
        remaining_items=items,
        n_samples=monte_carlo_samples,
        seed=seed,
        demand_mode=demand_mode,
    )

    # Compose payload
    payload: Dict[str, Any] = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "units": units,
        "completed_days": completed_days,
        "observed_cum_capacity": observed_cum_capacity,
        "demand_mode": demand_mode.value,
        "n_samples": int(monte_carlo_samples),
        "p_success": float(p_success),
        "workload_summary": workload_summary,
        "capacity_summary": capacity_summary,
        "plan_summary": {
            "n_developers": int(plan.n_devs),
            "sprint_days": int(plan.n_days),
            "developer_ids": list(plan.developer_ids),
        },
        "work_items_meta": work_items_meta or {},
    }

    if isinstance(work_items_meta, dict) and isinstance(work_items_meta.get("project_items"), list):
        payload["project_items"] = work_items_meta["project_items"]

    if bool(cfg.get("ai_uplift", {}).get("enabled", False)):
        result = forecaster.estimate_ai_uplift(
            completed_days=completed_days,
            observed_cum_capacity=observed_cum_capacity,
            remaining_items=items,
            n_samples=monte_carlo_samples,
            seed=seed,
            demand_mode=demand_mode,
        )
        payload["ai_uplift"] = {
            "p_ai": float(result.p_ai),
            "p_no_ai": float(result.p_no_ai),
            "p_no_ai_workload": float(result.p_no_ai_workload),
            "p_no_ai_capacity": float(result.p_no_ai_capacity),
            "delta_p_overall": float(result.delta_p_overall),
            "delta_p_from_workload": float(result.delta_p_from_workload),
            "delta_p_from_capacity": float(result.delta_p_from_capacity),
        }

    # Optional DM test
    try:
        dm_out = _maybe_compute_dm(cfg)
        if dm_out is not None:
            payload["forecast_accuracy"] = dm_out
    except Exception as e:
        payload["forecast_accuracy_error"] = str(e)

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    (outdir / "forecast.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")

    template_path = cfg.get("status_template_path") or (cfg.get("zenhub_status") or {}).get("template_path")
    if template_path:
        from sprint_forecast.status_rendering import render_status_markdown

        rendered = render_status_markdown(payload, template_path=str(template_path))
        payload["status_rendering"] = {"renderer": rendered.renderer, "template_path": rendered.template_path}
        (outdir / "forecast.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
        md = rendered.markdown
    else:
        md = _format_markdown_summary(payload)
    (outdir / "forecast.md").write_text(md, encoding="utf-8")

    # Optional Zenhub status update
    try:
        zh_update = _maybe_update_zenhub_status(cfg, md)
        if zh_update is not None:
            payload["zenhub_status_update"] = zh_update
            (outdir / "forecast.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
    except ZenhubError as e:
        payload["zenhub_status_update_error"] = str(e)
        (outdir / "forecast.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")

    print(f"Wrote {outdir / 'forecast.json'}")
    print(f"Wrote {outdir / 'forecast.md'}")


if __name__ == "__main__":
    main()
