"""
step6_generate_report.py
─────────────────────────────────────────────────────────────────────────────
Generates a Tier-3 descendant report by combining:
  1. Semantic search against the pgvector DB (war diary chunks)
  2. The prompt template from prompts/v2-tier3-report.md
  3. Claude (Anthropic) for the final report generation

Usage
─────
  # Default — uses the Bloggins placeholder soldier
  python scripts/step6_generate_report.py

  # Real soldier
  python scripts/step6_generate_report.py \
    --name "Pte. John Smith" \
    --unit "Calgary Highlanders" \
    --joined "mid September 1944" \
    --wounded "late October 1944" \
    --notes "Family knows he was in D Company. Was evacuated to England."

Output
──────
  reports/<soldier-slug>_<timestamp>_report.md
"""

import os
import sys
import argparse
import textwrap
from datetime import datetime
from pathlib import Path
from dotenv import load_dotenv

# ── Paths ─────────────────────────────────────────────────────────────────────
REPO_ROOT   = Path(__file__).resolve().parent.parent
PROMPT_FILE = REPO_ROOT / "prompts" / "v2-tier3-report.md"
REPORTS_DIR = REPO_ROOT / "reports"

# ── Add GoogleDocumentOCR to path so we can import db + embed_and_store ───────
sys.path.insert(0, str(REPO_ROOT / "GoogleDocumentOCR"))

load_dotenv()

# ── Search queries to give broad coverage of the service window ────────────────
SEARCH_QUERIES = [
    "Calgary Highlanders September 1944 operations movements France",
    "Calgary Highlanders October 1944 Scheldt fighting Holland",
    "Calgary Highlanders November 1944 casualties wounded Maas",
    "battalion advance attack objective company platoon",
    "reinforcements joined unit billets rest weather rations",
    "parades training administrative daily routine",
    "casualties killed wounded evacuated",
    "officers commanding lieutenant colonel major company commander",
    "civilian population towns villages liberated",
    "enemy German positions shelling mortars small arms fire",
]

CHUNKS_PER_QUERY = 5   # results per query
MAX_CONTEXT_CHARS = 80_000  # ~20k tokens — safe Claude window


def retrieve_diary_chunks(nationality: str = "canadian",
                           corpus: str = "war_diary_narrative") -> list[dict]:
    """
    Run multiple semantic searches and return a deduplicated, sorted list
    of diary chunks.
    """
    from embed_and_store import embed_texts
    from db import semantic_search

    seen_ids = set()
    chunks = []

    print(f"  Running {len(SEARCH_QUERIES)} semantic searches...")
    for query in SEARCH_QUERIES:
        embedding = embed_texts([query])[0]
        rows = semantic_search(
            embedding,
            top_k=CHUNKS_PER_QUERY,
            nationality=nationality,
            corpus=corpus,
        )
        for row in rows:
            filename, page_num, chunk_index, text, nat, corp, similarity = row
            key = (filename, page_num, chunk_index)
            if key not in seen_ids:
                seen_ids.add(key)
                chunks.append({
                    "filename":    filename,
                    "page_num":    page_num,
                    "chunk_index": chunk_index,
                    "text":        text,
                    "similarity":  similarity,
                })

    # Sort by filename then page so the context reads chronologically
    chunks.sort(key=lambda c: (c["filename"], c["page_num"], c["chunk_index"]))
    print(f"  Retrieved {len(chunks)} unique chunks after deduplication.")
    return chunks


def build_context_block(chunks: list[dict], max_chars: int = MAX_CONTEXT_CHARS) -> str:
    """
    Format chunks as a readable diary-extract block, trimmed to max_chars.
    """
    lines = []
    current_file = None
    char_count = 0

    for c in chunks:
        if c["filename"] != current_file:
            current_file = c["filename"]
            header = f"\n--- Source: {current_file} ---\n"
            lines.append(header)
            char_count += len(header)

        entry = f"[page {c['page_num']}] {c['text']}\n\n"
        if char_count + len(entry) > max_chars:
            lines.append("\n[Context truncated — additional sources available in DB]\n")
            break
        lines.append(entry)
        char_count += len(entry)

    return "".join(lines)


def build_customer_situation(name: str, unit: str, joined: str,
                              wounded: str, notes: str) -> str:
    knows = textwrap.dedent(f"""\
        The family knows:
        - Their relative, {name}, served with the {unit}.
        - They joined the battalion around {joined}.
        - They were wounded in action around {wounded}.
        - They were evacuated and did not return to the unit.
    """)

    doesnt_know = textwrap.dedent("""\
        The family does NOT know:
        - The exact date they joined or were wounded.
        - Specific actions they were personally involved in.
        - Their company, platoon, or section.
    """)

    if notes:
        knows += f"\nAdditional context provided by the family:\n{notes}\n"

    return f"## The customer's situation\n{knows}\n{doesnt_know}"


def load_prompt_template() -> str:
    return PROMPT_FILE.read_text(encoding="utf-8")


def assemble_full_prompt(template: str, customer_block: str,
                          context_block: str) -> str:
    """
    Replace the customer situation section in the template with the real
    soldier info, then append the diary context before the output instructions.
    """
    # Strip the hardcoded customer section and inject the real one
    start_marker = "## The customer's situation"
    end_marker   = "## What to produce"

    if start_marker in template and end_marker in template:
        before = template[:template.index(start_marker)]
        after  = template[template.index(end_marker):]
        template = before + customer_block + "\n\n" + after

    diary_section = (
        "## War diary source material\n"
        "The following excerpts are drawn from the pgvector database of "
        "OCR-processed war diaries. Use them as your primary source.\n\n"
        + context_block
        + "\n\n"
    )

    # Insert diary context just before "## What to produce"
    if end_marker in template:
        idx = template.index(end_marker)
        template = template[:idx] + diary_section + template[idx:]
    else:
        template += "\n\n" + diary_section

    return template


def generate_report_with_claude(prompt: str, soldier_name: str) -> str:
    import anthropic

    client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))

    print("  Sending to Claude... (this may take 30–60 seconds)")
    message = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=4096,
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
    )
    return message.content[0].text


def save_report(text: str, soldier_name: str) -> Path:
    REPORTS_DIR.mkdir(exist_ok=True)
    slug = soldier_name.lower().replace(" ", "_").replace(".", "").replace(",", "")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    out_path = REPORTS_DIR / f"{slug}_{timestamp}_report.md"
    out_path.write_text(text, encoding="utf-8")
    return out_path


def main():
    parser = argparse.ArgumentParser(
        description="Generate a Tier-3 descendant report from war diary pgvector DB"
    )
    parser.add_argument("--name",    default="Pte. Bill Bloggins",
                        help="Soldier's name (default: placeholder Bloggins)")
    parser.add_argument("--unit",    default="Calgary Highlanders",
                        help="Unit name")
    parser.add_argument("--joined",  default="mid September 1944",
                        help="Approximate join date (plain English)")
    parser.add_argument("--wounded", default="late October 1944",
                        help="Approximate wound/end date (plain English)")
    parser.add_argument("--notes",   default="",
                        help="Any extra family context (optional)")
    parser.add_argument("--nationality", default="canadian",
                        help="DB nationality filter (default: canadian)")
    parser.add_argument("--corpus",      default="war_diary_narrative",
                        help="DB corpus filter (default: war_diary_narrative)")
    parser.add_argument("--top-k", type=int, default=5,
                        help="Results per search query (default: 5)")
    args = parser.parse_args()

    print(f"\n=== Generating report for: {args.name} ===")
    print(f"    Unit    : {args.unit}")
    print(f"    Joined  : {args.joined}")
    print(f"    Wounded : {args.wounded}")
    if args.notes:
        print(f"    Notes   : {args.notes}")
    print()

    # 1. Retrieve diary chunks from pgvector
    print("[1/4] Retrieving diary chunks from pgvector DB...")
    chunks = retrieve_diary_chunks(
        nationality=args.nationality,
        corpus=args.corpus,
    )
    if not chunks:
        print("ERROR: No chunks found. Check DB connection and filters.")
        sys.exit(1)

    # 2. Build context block
    print("[2/4] Assembling context...")
    context_block = build_context_block(chunks)
    print(f"      Context: {len(context_block):,} characters from {len(chunks)} chunks")

    # 3. Assemble final prompt
    print("[3/4] Building prompt from v2-tier3-report.md...")
    template = load_prompt_template()
    customer_block = build_customer_situation(
        name=args.name, unit=args.unit,
        joined=args.joined, wounded=args.wounded,
        notes=args.notes,
    )
    full_prompt = assemble_full_prompt(template, customer_block, context_block)
    print(f"      Total prompt: {len(full_prompt):,} characters")

    # 4. Generate with Claude
    print("[4/4] Generating report with Claude...")
    report_text = generate_report_with_claude(full_prompt, args.name)

    # 5. Save
    out_path = save_report(report_text, args.name)
    print(f"\n✅ Report saved to: {out_path}")
    print("-" * 60)
    print(report_text[:500] + "...")


if __name__ == "__main__":
    main()

