"""Write all output artifacts for a backtest run."""
from __future__ import annotations

import csv
import datetime as dt
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq

from libs.backtest.domain import (
    BacktestConfig,
    DailyPortfolioState,
    ExperimentManifest,
    ExperimentResult,
    FilledTrade,
    MetricsBundle,
    OpenPosition,
)
from libs.common.config import get_settings
from libs.common.ids import sha256_checksum_str
from libs.common.logging import get_logger

logger = get_logger(__name__)


def create_run_directory(output_root: str | Path, run_id: str) -> Path:
    """Create the full run directory tree. Returns the run root path."""
    run_dir = Path(output_root) / run_id
    for subdir in ["logs", "metrics", "artifacts", "plots", "notes"]:
        (run_dir / subdir).mkdir(parents=True, exist_ok=True)
    logger.info("run_directory_created", path=str(run_dir))
    return run_dir


def write_metadata(
    run_dir: Path,
    run_id: str,
    started_at: dt.datetime,
    finished_at: dt.datetime,
    git_hash: str,
    total_trading_days: int,
    total_candidates_seen: int,
    total_orders_rejected: int,
    split_name: str | None = None,
) -> Path:
    """Write metadata.json."""
    meta = {
        "run_id": run_id,
        "started_at": started_at.isoformat(),
        "finished_at": finished_at.isoformat(),
        "elapsed_seconds": (finished_at - started_at).total_seconds(),
        "git_commit_hash": git_hash,
        "total_trading_days": total_trading_days,
        "total_candidates_seen": total_candidates_seen,
        "total_orders_rejected": total_orders_rejected,
    }
    if split_name is not None:
        meta["split_name"] = split_name
    out = run_dir / "metadata.json"
    out.write_text(json.dumps(meta, indent=2))
    return out


def write_metrics_summary(run_dir: Path, metrics: MetricsBundle) -> Path:
    """Write metrics/metrics_summary.json."""
    out = run_dir / "metrics" / "metrics_summary.json"
    out.write_text(metrics.model_dump_json(indent=2))
    return out


def write_snapshot_provenance(
    run_dir: Path,
    config: BacktestConfig,
) -> dict[str, str]:
    """Persist the snapshot manifest and a stable fingerprint for reproducibility checks."""
    snapshot_id = config.dataset_snapshot_id
    settings = get_settings()
    candidate_paths = [
        Path(settings.parquet_dir) / snapshot_id / "manifest.json",
        Path("data/datasets/snapshots") / snapshot_id / "manifest.json",
    ]
    manifest_path = next((path for path in candidate_paths if path.exists()), None)
    if manifest_path is None:
        return {}

    manifest_text = manifest_path.read_text()
    manifest_data = json.loads(manifest_text)
    manifest_snapshot_id = manifest_data.get("snapshot_id")
    manifest_output_dir = manifest_data.get("output_dir")
    snapshot_dir_name = manifest_path.parent.name
    snapshot_id_matches_manifest = manifest_snapshot_id == snapshot_dir_name
    output_dir_matches_manifest = (
        True if not manifest_output_dir else Path(str(manifest_output_dir)).name == snapshot_dir_name
    )
    if not snapshot_id_matches_manifest or not output_dir_matches_manifest:
        logger.warning(
            "snapshot_manifest_metadata_mismatch",
            dataset_snapshot_id=snapshot_id,
            source_manifest_path=str(manifest_path.resolve()),
            manifest_snapshot_id=manifest_snapshot_id,
            manifest_output_dir=manifest_output_dir,
            snapshot_dir_name=snapshot_dir_name,
        )
    manifest_copy = run_dir / "snapshot_manifest.json"
    manifest_copy.write_text(json.dumps(manifest_data, indent=2))

    fingerprint = {
        "dataset_snapshot_id": snapshot_id,
        "source_manifest_path": str(manifest_path.resolve()),
        "snapshot_dir_name": snapshot_dir_name,
        "manifest_snapshot_id": manifest_snapshot_id,
        "manifest_output_dir": manifest_output_dir,
        "snapshot_id_matches_manifest": snapshot_id_matches_manifest,
        "output_dir_matches_manifest": output_dir_matches_manifest,
        "manifest_sha256": sha256_checksum_str(manifest_text),
        "snapshot_created_at_utc": manifest_data.get("created_at_utc"),
        "code_commit_hash": manifest_data.get("code_commit_hash"),
        "feature_version": manifest_data.get("feature_version"),
        "parser_version": manifest_data.get("parser_version"),
        "label_version": manifest_data.get("label_version"),
        "row_counts": manifest_data.get("row_counts"),
        "total_rows": manifest_data.get("total_rows"),
    }
    fingerprint_path = run_dir / "snapshot_fingerprint.json"
    fingerprint_path.write_text(json.dumps(fingerprint, indent=2))
    return {
        "snapshot_manifest": str(manifest_copy),
        "snapshot_fingerprint": str(fingerprint_path),
    }


def write_trade_blotter(
    run_dir: Path,
    trades: list[FilledTrade],
) -> Path | None:
    """Write artifacts/trade_blotter.parquet."""
    if not trades:
        return None
    rows = [
        {
            "trade_id": t.trade_id,
            "position_id": t.position_id,
            "event_id": t.event_id,
            "symbol": t.symbol,
            "event_date": t.event_date.isoformat() if t.event_date else None,
            "event_type": t.event_type,
            "score": t.score,
            "timing_class": t.timing_class,
            "engine_id": t.engine_id,
            "entry_timing_policy": t.entry_timing_policy,
            "shadow_only": t.shadow_only,
            "parent_position_id": t.parent_position_id,
            "is_add_on": t.is_add_on,
            "entry_date": t.entry_date.isoformat(),
            "exit_date": t.exit_date.isoformat(),
            "entry_price": t.entry_price,
            "exit_price": t.exit_price,
            "exit_reason": t.exit_reason.value,
            "shares": t.shares,
            "commission": t.commission,
            "slippage_bps": t.slippage_bps,
            "gross_pnl": t.gross_pnl,
            "net_pnl": t.net_pnl,
            "pnl_pct": t.pnl_pct,
            "r_multiple": t.r_multiple,
            "holding_days": t.holding_days,
        }
        for t in trades
    ]
    out = run_dir / "artifacts" / "trade_blotter.parquet"
    _write_parquet(rows, out)
    return out


def write_daily_equity_curve(
    run_dir: Path,
    equity_curve: list[DailyPortfolioState],
) -> Path | None:
    """Write artifacts/daily_equity_curve.parquet."""
    if not equity_curve:
        return None
    rows = [
        {
            "date": s.date.isoformat(),
            "equity": s.equity,
            "cash_available": s.cash_available,
            "gross_exposure": s.gross_exposure,
            "net_exposure": s.net_exposure,
            "unrealized_pnl": s.unrealized_pnl,
            "realized_pnl": s.realized_pnl,
            "open_position_count": len(s.open_positions),
            "daily_new_risk_used": s.daily_new_risk_used,
            "peak_equity": s.peak_equity,
            "current_drawdown_pct": s.current_drawdown_pct,
        }
        for s in equity_curve
    ]
    out = run_dir / "artifacts" / "daily_equity_curve.parquet"
    _write_parquet(rows, out)
    return out


def write_position_timeline(
    run_dir: Path,
    trades: list[FilledTrade],
    open_positions: list[OpenPosition] | None = None,
) -> Path | None:
    """Write artifacts/position_timeline.parquet (one row per position)."""
    rows = []
    for t in trades:
        rows.append(
                {
                    "position_id": t.position_id,
                    "event_id": t.event_id,
                    "symbol": t.symbol,
                    "event_date": t.event_date.isoformat() if t.event_date else None,
                    "timing_class": t.timing_class,
                    "engine_id": t.engine_id,
                    "entry_timing_policy": t.entry_timing_policy,
                    "shadow_only": t.shadow_only,
                    "parent_position_id": t.parent_position_id,
                    "is_add_on": t.is_add_on,
                    "entry_date": t.entry_date.isoformat(),
                    "exit_date": t.exit_date.isoformat(),
                    "entry_price": t.entry_price,
                    "exit_price": t.exit_price,
                "exit_reason": t.exit_reason.value,
                "shares": t.shares,
                "net_pnl": t.net_pnl,
                "r_multiple": t.r_multiple,
                "holding_days": t.holding_days,
                "status": "closed",
            }
        )
    if open_positions:
        for p in open_positions:
            rows.append(
                {
                    "position_id": p.position_id,
                    "event_id": p.plan.candidate.event_id,
                    "symbol": p.plan.candidate.symbol,
                    "event_date": p.plan.event_date.isoformat() if p.plan.event_date else None,
                    "timing_class": p.plan.timing_class,
                    "engine_id": p.plan.engine_id,
                    "entry_timing_policy": p.plan.entry_timing_policy,
                    "shadow_only": p.plan.shadow_only,
                    "parent_position_id": p.parent_position_id,
                    "is_add_on": p.is_add_on,
                    "entry_date": p.entry_date.isoformat(),
                    "exit_date": None,
                    "entry_price": p.entry_price,
                    "exit_price": None,
                    "exit_reason": None,
                    "shares": p.shares_total,
                    "net_pnl": None,
                    "r_multiple": None,
                    "holding_days": p.days_held,
                    "status": p.status.value,
                }
            )
    if not rows:
        return None
    out = run_dir / "artifacts" / "position_timeline.parquet"
    _write_parquet(rows, out)
    return out


def write_attribution_by_event_type(
    run_dir: Path,
    trades: list[FilledTrade],
    candidate_map: dict[str, Any],
) -> Path:
    """Write metrics/attribution_by_event_type.csv."""
    bucket_data: dict[str, dict[str, float | int]] = defaultdict(
        lambda: {"count": 0, "wins": 0, "net_pnl": 0.0, "avg_r": 0.0, "_r_sum": 0.0}
    )
    for t in trades:
        cand = candidate_map.get(t.trade_id)
        et = getattr(cand, "event_type", "unknown") if cand else "unknown"
        d = bucket_data[et]
        d["count"] = int(d["count"]) + 1
        if t.net_pnl > 0:
            d["wins"] = int(d["wins"]) + 1
        d["net_pnl"] = float(d["net_pnl"]) + t.net_pnl
        d["_r_sum"] = float(d["_r_sum"]) + t.r_multiple

    out = run_dir / "metrics" / "attribution_by_event_type.csv"
    with open(out, "w", newline="") as f:
        writer = csv.DictWriter(
            f, fieldnames=["event_type", "count", "wins", "win_rate", "net_pnl", "avg_r"]
        )
        writer.writeheader()
        for et, d in sorted(bucket_data.items()):
            count = int(d["count"])
            wins = int(d["wins"])
            writer.writerow(
                {
                    "event_type": et,
                    "count": count,
                    "wins": wins,
                    "win_rate": wins / count if count > 0 else 0.0,
                    "net_pnl": round(float(d["net_pnl"]), 4),
                    "avg_r": round(float(d["_r_sum"]) / count if count > 0 else 0.0, 4),
                }
            )
    return out


def write_attribution_by_sector(
    run_dir: Path,
    trades: list[FilledTrade],
    candidate_map: dict[str, Any],
) -> Path:
    """Write metrics/attribution_by_sector.csv."""
    bucket_data: dict[str, dict[str, float | int]] = defaultdict(
        lambda: {"count": 0, "wins": 0, "net_pnl": 0.0, "_r_sum": 0.0}
    )
    for t in trades:
        cand = candidate_map.get(t.trade_id)
        sector = getattr(cand, "sector", "UNKNOWN") if cand else "UNKNOWN"
        d = bucket_data[sector]
        d["count"] = int(d["count"]) + 1
        if t.net_pnl > 0:
            d["wins"] = int(d["wins"]) + 1
        d["net_pnl"] = float(d["net_pnl"]) + t.net_pnl
        d["_r_sum"] = float(d["_r_sum"]) + t.r_multiple

    out = run_dir / "metrics" / "attribution_by_sector.csv"
    with open(out, "w", newline="") as f:
        writer = csv.DictWriter(
            f, fieldnames=["sector", "count", "wins", "win_rate", "net_pnl", "avg_r"]
        )
        writer.writeheader()
        for sector, d in sorted(bucket_data.items()):
            count = int(d["count"])
            wins = int(d["wins"])
            writer.writerow(
                {
                    "sector": sector,
                    "count": count,
                    "wins": wins,
                    "win_rate": wins / count if count > 0 else 0.0,
                    "net_pnl": round(float(d["net_pnl"]), 4),
                    "avg_r": round(float(d["_r_sum"]) / count if count > 0 else 0.0, 4),
                }
            )
    return out


def write_attribution_by_engine(
    run_dir: Path,
    trades: list[FilledTrade],
    per_engine_metrics: dict[str, dict[str, Any]] | None = None,
) -> Path:
    """Write metrics/attribution_by_engine.csv."""
    bucket_data: dict[str, dict[str, float | int | bool]] = defaultdict(
        lambda: {"count": 0, "wins": 0, "net_pnl": 0.0, "_r_sum": 0.0, "shadow_only": False}
    )
    for t in trades:
        engine_id = t.engine_id or "default"
        d = bucket_data[engine_id]
        d["count"] = int(d["count"]) + 1
        if t.net_pnl > 0:
            d["wins"] = int(d["wins"]) + 1
        d["net_pnl"] = float(d["net_pnl"]) + t.net_pnl
        d["_r_sum"] = float(d["_r_sum"]) + t.r_multiple
        d["shadow_only"] = bool(t.shadow_only)

    if per_engine_metrics:
        for engine_id, summary in per_engine_metrics.items():
            d = bucket_data.setdefault(
                engine_id,
                {"count": 0, "wins": 0, "net_pnl": 0.0, "_r_sum": 0.0, "shadow_only": False},
            )
            d["shadow_only"] = bool(summary.get("shadow_only", d["shadow_only"]))

    out = run_dir / "metrics" / "attribution_by_engine.csv"
    with open(out, "w", newline="") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "engine_id",
                "shadow_only",
                "count",
                "wins",
                "win_rate",
                "net_pnl",
                "avg_r",
            ],
        )
        writer.writeheader()
        for engine_id, d in sorted(bucket_data.items()):
            count = int(d["count"])
            wins = int(d["wins"])
            writer.writerow(
                {
                    "engine_id": engine_id,
                    "shadow_only": bool(d["shadow_only"]),
                    "count": count,
                    "wins": wins,
                    "win_rate": wins / count if count > 0 else 0.0,
                    "net_pnl": round(float(d["net_pnl"]), 4),
                    "avg_r": round(float(d["_r_sum"]) / count if count > 0 else 0.0, 4),
                }
            )
    return out


def write_score_bucket_report(
    run_dir: Path,
    score_bucket_hit_rate: dict[str, float],
    trades: list[FilledTrade],
    candidate_map: dict[str, Any],
) -> Path:
    """Write metrics/score_bucket_report.csv."""
    bucket_counts: dict[str, int] = defaultdict(int)
    for t in trades:
        cand = candidate_map.get(t.trade_id)
        bucket = getattr(cand, "score_bucket", "unknown") if cand else "unknown"
        bucket_counts[bucket] += 1

    out = run_dir / "metrics" / "score_bucket_report.csv"
    with open(out, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["score_bucket", "trade_count", "win_rate"])
        writer.writeheader()
        for bucket in sorted(set(list(score_bucket_hit_rate.keys()) + list(bucket_counts.keys()))):
            writer.writerow(
                {
                    "score_bucket": bucket,
                    "trade_count": bucket_counts.get(bucket, 0),
                    "win_rate": round(score_bucket_hit_rate.get(bucket, 0.0), 4),
                }
            )
    return out


def write_plots(run_dir: Path, generate: bool = False) -> Path:
    """Create plots directory. generate=True logs a warning (matplotlib not available)."""
    plots_dir = run_dir / "plots"
    plots_dir.mkdir(exist_ok=True)
    if generate:
        logger.warning(
            "plots_not_implemented",
            message="generate_plots=True is a no-op; matplotlib is not in dependencies.",
        )
    return plots_dir


def write_run_notes(run_dir: Path, notes: str = "") -> Path:
    """Write notes/run_notes.md."""
    out = run_dir / "notes" / "run_notes.md"
    out.write_text(notes or "# Run Notes\n\n_No notes provided._\n")
    return out


def write_per_engine_metrics(
    run_dir: Path,
    per_engine_metrics: dict[str, dict[str, Any]],
) -> Path:
    """Write metrics/per_engine_metrics.json."""
    out = run_dir / "metrics" / "per_engine_metrics.json"
    out.write_text(json.dumps(per_engine_metrics, indent=2, default=str))
    return out


def write_all_artifacts(
    run_dir: Path,
    run_id: str,
    manifest: ExperimentManifest,
    config: BacktestConfig,
    metrics: MetricsBundle,
    trades: list[FilledTrade],
    equity_curve: list[DailyPortfolioState],
    open_positions: list[OpenPosition],
    candidate_map: dict[str, Any],
    started_at: dt.datetime,
    finished_at: dt.datetime,
    git_hash: str,
    total_trading_days: int,
    total_candidates_seen: int,
    total_orders_rejected: int,
    split_name: str | None = None,
    per_engine_metrics: dict[str, dict[str, Any]] | None = None,
) -> dict[str, str]:
    """Write all output files. Returns mapping of artifact_name → file_path."""
    from libs.backtest.manifests import save_manifest, save_resolved_config

    paths: dict[str, str] = {}

    # Core files
    paths["manifest"] = str(save_manifest(manifest, run_dir))
    paths["resolved_config"] = str(save_resolved_config(config, run_dir))
    paths.update(write_snapshot_provenance(run_dir, config))
    paths["metadata"] = str(
        write_metadata(
            run_dir, run_id, started_at, finished_at, git_hash,
            total_trading_days, total_candidates_seen, total_orders_rejected,
            split_name=split_name,
        )
    )

    # Metrics
    if config.reporting.write_metrics_summary:
        paths["metrics_summary"] = str(write_metrics_summary(run_dir, metrics))
    paths["attribution_by_event_type"] = str(
        write_attribution_by_event_type(run_dir, trades, candidate_map)
    )
    paths["attribution_by_sector"] = str(
        write_attribution_by_sector(run_dir, trades, candidate_map)
    )
    paths["attribution_by_engine"] = str(
        write_attribution_by_engine(run_dir, trades, per_engine_metrics)
    )
    paths["score_bucket_report"] = str(
        write_score_bucket_report(
            run_dir, metrics.score_bucket_hit_rate, trades, candidate_map
        )
    )
    if per_engine_metrics:
        paths["per_engine_metrics"] = str(write_per_engine_metrics(run_dir, per_engine_metrics))

    # Trade data
    if config.reporting.write_trade_blotter:
        p = write_trade_blotter(run_dir, trades)
        if p:
            paths["trade_blotter"] = str(p)

    if config.reporting.write_equity_curve:
        p = write_daily_equity_curve(run_dir, equity_curve)
        if p:
            paths["daily_equity_curve"] = str(p)

    p = write_position_timeline(run_dir, trades, open_positions)
    if p:
        paths["position_timeline"] = str(p)

    # Plots (no-op)
    paths["plots_dir"] = str(write_plots(run_dir, config.reporting.generate_plots))

    # Notes
    paths["run_notes"] = str(write_run_notes(run_dir, manifest.notes or ""))

    return paths


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _write_parquet(rows: list[dict[str, Any]], path: Path) -> None:
    """Write list of row dicts to Parquet."""
    if not rows:
        return
    keys = list(rows[0].keys())
    columns: dict[str, list[Any]] = {k: [] for k in keys}
    for row in rows:
        for k in keys:
            columns[k].append(row.get(k))
    table = pa.table({k: pa.array(v) for k, v in columns.items()})
    pq.write_table(table, str(path))
    logger.debug("parquet_written", path=str(path), rows=len(rows))