"""Evidence extraction and analysis — Sonnet with web search, structured output."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from superpitch.claude import SONNET, call, text_from


# ── Shared helpers ────────────────────────────────────────────────────────────

def _attachment_blocks(attachments: list[dict]) -> tuple[list[dict], int]:
    """Build Anthropic content blocks for PDFs and images. Returns (blocks, pdf_count)."""
    blocks: list[dict] = []
    pdf_count = 0
    for att in attachments:
        mt = att.get("content_type", "")
        fn = att.get("filename", "").lower()
        if (mt.startswith("application/pdf") or fn.endswith(".pdf")) and att.get("data"):
            if pdf_count >= 3:
                continue
            blocks.append({
                "type": "document",
                "source": {"type": "base64", "media_type": "application/pdf", "data": att["data"]},
                "title": att.get("filename", "pitch.pdf"),
            })
            pdf_count += 1
        elif mt.startswith("image/") and att.get("data"):
            blocks.append({
                "type": "image",
                "source": {"type": "base64", "media_type": mt, "data": att["data"]},
            })
    return blocks, pdf_count


def specialist_features(packet: dict) -> dict:
    """Extract the feature dict used by the Iris specialist from a packet."""
    return {
        "stage":      packet.get("stage", ""),
        "sector":     packet.get("sector", ""),
        "geography":  packet.get("geography", ""),
        "ask":        packet.get("ask", ""),
        "traction":   packet.get("traction", ""),
        "team":       packet.get("team", ""),
        "thesis_fit": packet.get("thesis_fit", ""),
        "missing":    ", ".join(packet.get("missing") or []),
        "confidence": packet.get("confidence", 0.0),
    }


# ── Schemas ──────────────────────────────────────────────────────────────────

_EVIDENCE_SCHEMA = {
    "type": "json_schema",
    "schema": {
        "type": "object",
        "properties": {
            "company":       {"type": "string"},
            "one_liner":     {"type": "string"},
            "founder_name":  {"type": "string"},
            "founder_email": {"type": "string"},
            "stage":         {"type": "string"},
            "sector":        {"type": "string"},
            "geography":     {"type": "string"},
            "ask":           {"type": "string"},
            "traction":      {"type": "string"},
            "team":          {"type": "string"},
            "business_model":{"type": "string"},
            "thesis_fit":    {"type": "string"},
            "research":      {"type": "array", "items": {"type": "object", "properties": {"finding": {"type": "string"}, "source": {"type": "string"}}, "required": ["finding", "source"]}},
            "missing":       {"type": "array", "items": {"type": "string"}},
            "follow_up":     {"type": "array", "items": {"type": "string"}},
            "action":        {"type": "string"},
            "confidence":    {"type": "number"},
            "action_reason": {"type": "string"},
        },
        "required": ["company", "one_liner", "action", "confidence", "action_reason"],
        "additionalProperties": False,
    },
}

_FINAL_SCHEMA = {
    "type": "json_schema",
    "schema": {
        "type": "object",
        "properties": {
            "questions":        {"type": "array", "items": {"type": "string"}},
            "verdict":          {"type": "string"},
            "verdict_reasoning":{"type": "string"},
            "thank_you_subject":{"type": "string"},
            "thank_you_body":   {"type": "string"},
            "action":           {"type": "string"},
            "confidence":       {"type": "number"},
            "action_reason":    {"type": "string"},
        },
        "required": ["questions", "verdict", "verdict_reasoning", "thank_you_subject", "thank_you_body", "action", "confidence", "action_reason"],
        "additionalProperties": False,
    },
}


# ── Specialist triage ────────────────────────────────────────────────────────

def _triage_with_specialist(packet: dict, thesis: str = "") -> dict:
    """Override action with the bundled Iris specialist. Falls back silently."""
    try:
        from specialist.ask import ask  # type: ignore[import]
    except ImportError:
        return packet

    try:
        result = ask(specialist_features(packet), thesis=thesis)
    except Exception:
        return packet

    text = result if isinstance(result, str) else str(result)
    for action in ("ESCALATE", "REVIEW", "PASS"):
        if action in text.upper():
            packet["action"] = action
            packet["action_reason"] = text.split(".", 1)[0].strip()
            break
    return packet


# ── Extraction ───────────────────────────────────────────────────────────────

_EXTRACT_SYSTEM = (
    "You are an analyst helping an early-stage investor screen inbound pitches. "
    "Extract structured facts from the email and any attachments, research the company "
    "and founder online, then return a JSON evidence packet. "
    "Be factual and concise. Cite sources for research findings. "
    "Omit fields you cannot determine rather than guessing. "
    "If anything important is unclear or missing, list specific clarification questions "
    "in follow_up[] — these will be emailed directly to the founder."
)

_EXTRACT_PROMPT = """\
Investor thesis: {thesis}

Inbound pitch email —
Subject: {subject}
From: {sender}

Body:
{body}

{attachments_note}

Search the web for the company and founder, then return the evidence packet.

Action guide:
- ESCALATE: strong team + clear traction + thesis fit. Investor should look now.
- REVIEW: promising but incomplete or uncertain. Worth a closer look.
- PASS: weak fit, no key signals, or clearly outside thesis."""


def extract_evidence_packet(raw_email: dict, thesis: str, api_key: str) -> dict | None:
    """Use Sonnet with web search + PDF vision to extract a structured evidence packet."""
    blocks, pdf_count = _attachment_blocks(raw_email.get("attachments", []))
    attachments_note = (
        f"{pdf_count} PDF attachment(s) provided above — read them carefully."
        if pdf_count else "No attachments."
    )
    blocks.append({"type": "text", "text": _EXTRACT_PROMPT.format(
        thesis=thesis or "early-stage technology startups",
        subject=raw_email.get("subject", ""),
        sender=raw_email.get("from", ""),
        body=raw_email.get("body_text", "")[:4000],
        attachments_note=attachments_note,
    )})

    resp = call(
        SONNET, [{"role": "user", "content": blocks}],
        api_key=api_key, tools=[{"type": "web_search_20260209", "name": "web_search"}],
        output_config={"format": _EVIDENCE_SCHEMA},
        max_tokens=2048, system=_EXTRACT_SYSTEM,
    )
    if resp.get("error"):
        return None
    try:
        packet = json.loads(text_from(resp).strip())
    except Exception:
        return None
    if not packet:
        return None
    if not packet.get("follow_up"):
        packet = finalize_packet(packet, thesis, api_key)
    return packet


# ── Final pass ───────────────────────────────────────────────────────────────

_FINAL_SYSTEM = (
    "You are an analyst producing the final assessment of a startup pitch for an "
    "early-stage investor. The evidence packet is complete — all clarifications have "
    "been gathered. Now produce your final analysis."
)

_FINAL_PROMPT = """\
Investor thesis: {thesis}

Complete evidence packet:
{packet}

Based on the full evidence packet above:
1. Write 3-5 sharp questions the investor should ask the founder on a first call.
2. Give your verdict — is this a good investment opportunity? Consider team, traction, market, thesis fit, and risks.
3. Decide the final triage action (ESCALATE / REVIEW / PASS) with confidence.
4. Write a brief, professional thank-you email to the founder (subject line + body). \
Keep it warm but concise — thank them for the information, say you have what you need for now."""


def finalize_packet(packet: dict, thesis: str, api_key: str) -> dict:
    """Run the final pass: add questions[], verdict, thank-you, and revised action/confidence."""
    resp = call(
        SONNET,
        [{"role": "user", "content": _FINAL_PROMPT.format(
            thesis=thesis or "early-stage technology startups",
            packet=json.dumps(packet, indent=2)[:4000],
        )}],
        api_key=api_key, output_config={"format": _FINAL_SCHEMA},
        max_tokens=1024, system=_FINAL_SYSTEM,
    )
    if resp.get("error"):
        return packet

    try:
        final = json.loads(text_from(resp).strip())
    except Exception:
        return packet

    for key in ("questions", "verdict", "verdict_reasoning", "thank_you_subject",
                "thank_you_body", "action", "confidence", "action_reason"):
        if final.get(key) is not None:
            packet[key] = final[key]

    return _triage_with_specialist(packet, thesis=thesis)


# ── File-based one-off mode ──────────────────────────────────────────────────

def process_file(path: str, thesis: str, api_key: str) -> dict | None:
    """Process a local file (PDF or text) for one-off analysis."""
    import base64 as b64

    p = Path(path)
    raw_email: dict[str, Any] = {"subject": p.name, "from": "local file", "body_text": "", "attachments": []}
    if p.suffix.lower() == ".pdf":
        raw_email["attachments"] = [{
            "filename": p.name, "content_type": "application/pdf",
            "data": b64.b64encode(p.read_bytes()).decode(),
        }]
    else:
        raw_email["body_text"] = p.read_text(encoding="utf-8", errors="ignore")

    return extract_evidence_packet(raw_email, thesis, api_key)
