from __future__ import annotations

import argparse
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
RUNS_DIR = ROOT / "runs"
HINDSIGHT_SITE = ROOT / ".venv-hindsight" / "lib" / "python3.12" / "site-packages"
if HINDSIGHT_SITE.exists():
    sys.path.insert(0, str(HINDSIGHT_SITE))

os.environ.setdefault("HOME", "/tmp/hindsight-benchmark-import-home")
os.environ.setdefault("HINDSIGHT_API_LLM_PROVIDER", "none")
os.environ.setdefault("HINDSIGHT_API_LOG_LEVEL", "warning")

from hindsight import HindsightEmbedded  # noqa: E402

from run_sibyl_365d_500c_category_benchmark import (  # noqa: E402
    MARKER_DAYS,
    MILESTONE_DAYS,
    ROLES,
    START_DATE,
    STAKEHOLDERS_PER_COMPANY,
    TIMELINE_DAYS,
    TOPICS,
    approx_tokens,
    day_date,
    display,
    marker,
    region,
    score_context,
    segment,
    slug,
    spread_indices,
    topic,
    write_json,
)


DEFAULT_RUN_ID = "run-2026-06-07-hindsight-365d-500c-50q-category-retrieval"
DEFAULT_DATA_ROOT = f"/tmp/{DEFAULT_RUN_ID}"
RUN_SCOPE = "365d-500c-50q-hindsight-baseline"


def iso_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def flatten_memory(value: Any, *, prefix: str = "") -> list[str]:
    if isinstance(value, dict):
        tokens: list[str] = []
        for key, item in value.items():
            nested_prefix = f"{prefix}.{key}" if prefix else str(key)
            tokens.extend(flatten_memory(item, prefix=nested_prefix))
        return tokens
    if isinstance(value, list | tuple):
        tokens: list[str] = []
        for index, item in enumerate(value):
            tokens.extend(flatten_memory(item, prefix=f"{prefix}.{index}"))
        return tokens
    if value is None:
        return []
    return [f"{prefix}={value}"]


def make_memory(data: dict[str, Any]) -> str:
    return " ".join(flatten_memory(data))


def company_key(index: int) -> str:
    return f"run365-{slug(index)}"


def person_key(index: int, role_index: int) -> str:
    return f"run365-person-{index:03d}-{role_index:02d}"


def relation_key(index: int, role_index: int) -> str:
    return f"run365-rel-{index:03d}-{role_index:02d}"


def company_records(index: int, *, include_state_snapshots: bool) -> list[str]:
    company_slug = slug(index)
    company_display = display(index)
    records: list[str] = []
    records.append(
        make_memory(
            {
                "tier": "entity",
                "category": "company",
                "key": company_key(index),
                "body": {
                    "display_name": company_display,
                    "segment": segment(index),
                    "region": region(index),
                    "timeline_start": START_DATE.isoformat(),
                    "timeline_days": TIMELINE_DAYS,
                    "stakeholder_count": STAKEHOLDERS_PER_COMPANY,
                    "context_stats": {
                        "simulated_days": TIMELINE_DAYS,
                        "topic_cycle_size": len(TOPICS),
                        "milestone_count": len(MILESTONE_DAYS),
                    },
                },
            }
        )
    )
    for role_index, role in enumerate(ROLES, start=1):
        person = person_key(index, role_index)
        records.append(
            make_memory(
                {
                    "tier": "entity",
                    "category": "person",
                    "key": person,
                    "body": {
                        "display_name": f"{company_display} Contact {role_index}",
                        "role": role,
                        "company": company_key(index),
                        "company_display": company_display,
                    },
                }
            )
        )
        records.append(
            make_memory(
                {
                    "tier": "entity",
                    "category": "relationship",
                    "key": relation_key(index, role_index),
                    "body": {
                        "source": person,
                        "target": company_key(index),
                        "type": "works_on_account",
                        "role": role,
                        "description": f"{person} is the {role} for {company_display}.",
                    },
                }
            )
        )
    for day_index in range(1, TIMELINE_DAYS + 1):
        event_marker = marker(index, day_index)
        event_topic = topic(day_index)
        event_date = day_date(day_index)
        records.append(
            make_memory(
                {
                    "tier": "journal",
                    "category": "chronology",
                    "key": f"run365-{company_slug}-day-{day_index:03d}",
                    "body": {
                        "kind": "scale365_chronology_daily_update",
                        "company": company_display,
                        "company_slug": company_slug,
                        "day_index": day_index,
                        "date": event_date,
                        "topic": event_topic,
                        "marker": event_marker,
                        "summary": f"{company_display} day {day_index:03d}: {event_topic} update recorded.",
                    },
                }
            )
        )
        if day_index in MILESTONE_DAYS:
            records.append(
                make_memory(
                    {
                        "tier": "entity",
                        "category": "timeline_milestone",
                        "key": f"run365-{company_slug}-day-{day_index:03d}",
                        "body": {
                            "company": company_key(index),
                            "company_display": company_display,
                            "day_index": day_index,
                            "date": event_date,
                            "topic": event_topic,
                            "marker": event_marker,
                            "decision": f"{company_display} checkpoint day {day_index:03d} accepted for {event_topic}.",
                        },
                    }
                )
            )
        if include_state_snapshots and (day_index % 10 == 0 or day_index == TIMELINE_DAYS):
            records.append(
                make_memory(
                    {
                        "tier": "state",
                        "key": f"run365-current-status-{company_slug}",
                        "snapshot_day_index": day_index,
                        "body": {
                            "company": company_display,
                            "last_day_index": day_index,
                            "last_date": event_date,
                            "last_marker": event_marker,
                            "status": "active",
                        },
                    }
                )
            )
    if not include_state_snapshots:
        records.append(
            make_memory(
                {
                    "tier": "state",
                    "key": f"run365-current-status-{company_slug}",
                    "body": {
                        "company": company_display,
                        "last_day_index": TIMELINE_DAYS,
                        "last_date": day_date(TIMELINE_DAYS),
                        "last_marker": marker(index, TIMELINE_DAYS),
                        "status": "active",
                    },
                }
            )
        )
    return records


def build_questions(company_count: int, questions_per_category: int) -> list[dict[str, Any]]:
    cases: list[dict[str, Any]] = []
    for index in spread_indices(questions_per_category, company_count, offset=0):
        cases.append(
            {
                "category": "status",
                "id": f"status_company_{index:03d}",
                "query": f"What is the current status for {display(index)} after the 365 day chronology?",
                "retrieval_query": f"{display(index)} current status day 365",
                "expected_contains": [display(index), str(TIMELINE_DAYS), day_date(TIMELINE_DAYS), "active"],
            }
        )
    for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=7)):
        day_index = MILESTONE_DAYS[position % len(MILESTONE_DAYS)]
        cases.append(
            {
                "category": "milestone",
                "id": f"milestone_company_{index:03d}_day_{day_index:03d}",
                "query": f"What was the {display(index)} day {day_index} milestone accepted for?",
                "retrieval_query": f"{display(index)} day {day_index} milestone",
                "expected_contains": [display(index), str(day_index), day_date(day_index), topic(day_index)],
            }
        )
    for index in spread_indices(questions_per_category, company_count, offset=13):
        cases.append(
            {
                "category": "context_stat",
                "id": f"context_company_{index:03d}",
                "query": f"What segment and region context is stored for {display(index)}?",
                "retrieval_query": display(index),
                "expected_contains": [display(index), segment(index), region(index), str(TIMELINE_DAYS)],
            }
        )
    for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=19)):
        role_index = (position % len(ROLES)) + 1
        cases.append(
            {
                "category": "role",
                "id": f"role_company_{index:03d}_{role_index}",
                "query": f"Who is the {ROLES[role_index - 1]} for {display(index)}?",
                "retrieval_query": f"{display(index)} {ROLES[role_index - 1]}",
                "expected_contains": [f"{display(index)} Contact {role_index}", ROLES[role_index - 1]],
            }
        )
    for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=29)):
        day_index = MARKER_DAYS[position % len(MARKER_DAYS)]
        cases.append(
            {
                "category": "marker",
                "id": f"marker_company_{index:03d}_day_{day_index:03d}",
                "query": f"What marker is associated with {display(index)} day {day_index}?",
                "retrieval_query": marker(index, day_index),
                "expected_contains": [marker(index, day_index)],
            }
        )
    for position, index in enumerate(spread_indices(questions_per_category, company_count, offset=37)):
        day_index = 1 + ((position * 17) % TIMELINE_DAYS)
        cases.append(
            {
                "category": "temporal_topic",
                "id": f"topic_company_{index:03d}_day_{day_index:03d}",
                "query": f"What topic was recorded for {display(index)} on day {day_index}?",
                "retrieval_query": f"{display(index)} day {day_index} topic",
                "expected_contains": [display(index), str(day_index), day_date(day_index), topic(day_index)],
            }
        )
    for index in range(1, questions_per_category + 1):
        name = f"Phantom Prospect {index:03d}"
        safe = name.lower().replace(" ", "_")
        cases.append(
            {
                "category": "negative_trap",
                "id": f"negative_{safe}",
                "query": f"What is the current status for {name} after the 365 day chronology?",
                "retrieval_query": f"{name} current status day 365",
                "expected_unknown": True,
                "forbidden_contains": ["active", str(TIMELINE_DAYS), day_date(TIMELINE_DAYS)],
            }
        )
    return cases


def ingest(client: HindsightEmbedded, bank_id: str, *, company_count: int, batch_size: int, include_state_snapshots: bool, progress: bool) -> dict[str, Any]:
    started = time.perf_counter()
    records_added = 0
    for index in range(1, company_count + 1):
        records = company_records(index, include_state_snapshots=include_state_snapshots)
        for offset in range(0, len(records), batch_size):
            batch = [
                {
                    "content": content,
                    "metadata": {"benchmark": "sibyl-365d-500c-category", "company": display(index), "company_index": str(index)},
                    "tags": [RUN_SCOPE, f"company-{index:03d}"],
                }
                for content in records[offset : offset + batch_size]
            ]
            client.retain_batch(bank_id=bank_id, items=batch, retain_async=False)
            records_added += len(batch)
        if progress and (index == 1 or index % 10 == 0):
            print(f"[ingest] company {index}/{company_count} records={records_added} elapsed={time.perf_counter() - started:.1f}s", file=sys.stderr, flush=True)
    return {"mode": "hindsight_embedded_provider_none_retain_batch", "records_added": records_added, "elapsed_seconds": round(time.perf_counter() - started, 3)}


def result_row(item: dict[str, Any]) -> dict[str, Any]:
    text = str(item.get("text") or "")
    return {
        "tier": item.get("type") or "hindsight",
        "key": item.get("id"),
        "category": (item.get("metadata") or {}).get("company"),
        "body": text,
        "snippet": text[:500],
        "document_id": item.get("document_id"),
        "chunk_id": item.get("chunk_id"),
        "tags": item.get("tags"),
    }


def search_context(client: HindsightEmbedded, bank_id: str, case: dict[str, Any], *, max_tokens: int, include_chunks: bool) -> list[dict[str, Any]]:
    query = str(case.get("retrieval_query") or case.get("query") or "")
    response = client.recall(
        bank_id=bank_id,
        query=query,
        max_tokens=max_tokens,
        budget="low",
        include_chunks=include_chunks,
        tags=[RUN_SCOPE],
        tags_match="all_strict",
    )
    data = response.model_dump()
    rows = [result_row(item) for item in data.get("results") or []]
    chunks = data.get("chunks") or {}
    for chunk_id, chunk in chunks.items():
        text = str((chunk or {}).get("text") or "")
        rows.append(
            {
                "tier": "chunk",
                "key": chunk_id,
                "category": None,
                "body": text,
                "snippet": text[:500],
                "document_id": None,
                "chunk_id": chunk_id,
                "tags": None,
            }
        )
    return rows


def summarize_questions(questions: list[dict[str, Any]]) -> dict[str, Any]:
    summary: dict[str, Any] = {}
    for item in questions:
        category = item["category"]
        bucket = summary.setdefault(category, {"passed": 0, "total": 0})
        bucket["total"] += 1
        bucket["passed"] += int(bool(item["retrieval_score"]["passed"]))
    return summary


def build_markdown(raw: dict[str, Any]) -> str:
    rows = "\n".join(
        f"| {category} | {row['passed']} / {row['total']} |"
        for category, row in raw["summary"]["retrieval_by_category"].items()
    )
    failures = [
        {
            "id": item["id"],
            "category": item["category"],
            "query": item["query"],
            "expected_contains": item.get("expected_contains", []),
            "forbidden_contains": item.get("forbidden_contains", []),
            "retrieval_score": item["retrieval_score"],
            "context_sample": item["context"][:3],
        }
        for item in raw["questions"]
        if not item["retrieval_score"]["passed"]
    ]
    return f"""# Hindsight 365d 500c Category Baseline

Run: `{raw["run_id"]}`

## Scope

- System: Hindsight embedded local daemon.
- Corpus: {raw["dataset"]["company_count"]} companies, {raw["dataset"]["stakeholder_count"]} stakeholders, {raw["dataset"]["timeline_days"]} simulated days.
- Stored memories: {raw["ingest"]["records_added"]}.
- Questions: {raw["summary"]["question_total"]}, {raw["dataset"]["questions_per_category"]} per category.
- Mode: `{raw["ingest"]["mode"]}`.
- LLM provider: `{raw["hindsight"]["llm_provider"]}`.
- Observations/consolidation: disabled.
- Retrieval: Hindsight `recall`, max_tokens={raw["hindsight"]["max_tokens"]}, budget=low, include_chunks={raw["hindsight"]["include_chunks"]}.
- Answer mode: skipped.

## Summary

| Metric | Value |
| --- | ---: |
| Retrieval passed | {raw["summary"]["retrieval_passed"]} / {raw["summary"]["question_total"]} |
| Avg context rows | {raw["summary"]["avg_context_rows"]:.2f} |
| Avg context tokens | {raw["summary"]["avg_context_tokens"]:.2f} |
| Ingest elapsed seconds | {raw["ingest"]["elapsed_seconds"]} |
| Retrieval elapsed seconds | {raw["summary"]["retrieval_elapsed_seconds"]} |
| Model/API cost USD | 0.000000 |

## Retrieval By Category

| Category | Retrieval passed |
| --- | ---: |
{rows}

## Failures

```json
{json.dumps(failures[:25], indent=2, sort_keys=True, default=str)}
```
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run Hindsight embedded baseline on the Sibyl 365d/500-company category benchmark.")
    parser.add_argument("--run-id", default=DEFAULT_RUN_ID)
    parser.add_argument("--data-root", default=DEFAULT_DATA_ROOT)
    parser.add_argument("--company-limit", type=int, default=500)
    parser.add_argument("--questions-per-category", type=int, default=50)
    parser.add_argument("--question-limit", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=64)
    parser.add_argument("--max-tokens", type=int, default=12000)
    parser.add_argument("--skip-ingest", action="store_true")
    parser.add_argument("--include-state-snapshots", action="store_true")
    parser.add_argument("--no-chunks", action="store_true")
    parser.add_argument("--stop-daemon", action="store_true")
    parser.add_argument("--progress", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.company_limit < 1 or args.company_limit > 500:
        raise SystemExit("--company-limit must be between 1 and 500")
    if args.questions_per_category < 1:
        raise SystemExit("--questions-per-category must be positive")

    data_root = Path(args.data_root)
    data_root.mkdir(parents=True, exist_ok=True)
    os.environ["HOME"] = str(data_root / "home")
    os.environ["HINDSIGHT_API_LLM_PROVIDER"] = "none"
    os.environ["HINDSIGHT_API_LOG_LEVEL"] = "warning"

    profile = args.run_id.replace("run-", "").replace(".", "-")
    bank_id = args.run_id.replace("_", "-")
    client = HindsightEmbedded(profile=profile, llm_provider="none", idle_timeout=300, log_level="warning")
    try:
        raw: dict[str, Any] = {
            "run_id": args.run_id,
            "ts_started": iso_now(),
            "source": str(data_root),
            "hindsight": {
                "version": "0.7.2",
                "run_scope": RUN_SCOPE,
                "profile": profile,
                "bank_id": bank_id,
                "url": client.url,
                "llm_provider": "none",
                "max_tokens": args.max_tokens,
                "include_chunks": not args.no_chunks,
            },
        }
        client.create_bank(bank_id=bank_id, enable_observations=False, retain_extraction_mode="concise")
        if args.skip_ingest:
            ingest_result = {"mode": "skipped", "records_added": None, "elapsed_seconds": 0}
        else:
            ingest_result = ingest(
                client,
                bank_id,
                company_count=args.company_limit,
                batch_size=args.batch_size,
                include_state_snapshots=args.include_state_snapshots,
                progress=args.progress,
            )

        questions = build_questions(args.company_limit, args.questions_per_category)
        if args.question_limit is not None:
            questions = questions[: args.question_limit]

        started_questions = time.perf_counter()
        question_rows: list[dict[str, Any]] = []
        for position, case in enumerate(questions, start=1):
            context = search_context(client, bank_id, case, max_tokens=args.max_tokens, include_chunks=not args.no_chunks)
            row = {
                **case,
                "context_count": len(context),
                "context_tokens_estimate": approx_tokens(context),
                "context": context,
                "retrieval_score": score_context(case, context),
            }
            question_rows.append(row)
            if args.progress and (position == 1 or position % 25 == 0):
                print(f"[questions] {position}/{len(questions)} elapsed={time.perf_counter() - started_questions:.1f}s", file=sys.stderr, flush=True)

        retrieval_passed = sum(1 for item in question_rows if item["retrieval_score"]["passed"])
        raw["ts_completed"] = iso_now()
        raw["ingest"] = ingest_result
        raw["dataset"] = {
            "company_count": args.company_limit,
            "stakeholder_count": args.company_limit * STAKEHOLDERS_PER_COMPANY,
            "timeline_days": TIMELINE_DAYS,
            "questions_per_category": args.questions_per_category,
            "include_state_snapshots": args.include_state_snapshots,
        }
        raw["questions"] = question_rows
        raw["summary"] = {
            "question_total": len(question_rows),
            "retrieval_passed": retrieval_passed,
            "retrieval_failed": len(question_rows) - retrieval_passed,
            "retrieval_by_category": summarize_questions(question_rows),
            "avg_context_rows": sum(item["context_count"] for item in question_rows) / len(question_rows),
            "avg_context_tokens": sum(item["context_tokens_estimate"] for item in question_rows) / len(question_rows),
            "retrieval_elapsed_seconds": round(time.perf_counter() - started_questions, 3),
            "estimated_cost": {"estimated_total_usd": 0.0},
        }
        raw["status"] = "PASS" if retrieval_passed == len(question_rows) else "FAIL"

        RUNS_DIR.mkdir(parents=True, exist_ok=True)
        raw_path = RUNS_DIR / f"{args.run_id}.raw_result.json"
        md_path = RUNS_DIR / f"{args.run_id}.md"
        write_json(raw_path, raw)
        md_path.write_text(build_markdown(raw), encoding="utf8")
        print(json.dumps({"status": raw["status"], "summary": raw["summary"], "raw": str(raw_path), "md": str(md_path)}, indent=2))
    finally:
        client.close(stop_daemon=args.stop_daemon)


if __name__ == "__main__":
    main()