"""Email processing pipeline — classify with Haiku, extract and update with Sonnet."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from superpitch.claude import HAIKU, SONNET, call, parse_json, text_from


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

def _attachment_blocks(attachments: list[dict], default_title: str = "attachment") -> 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", default_title),
            })
            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


# ── Classification ─────────────────────────────────────────────────────────────

_CLASSIFY_PROMPT = """\
Classify this email as exactly one of: pitch, noise, unsure

PITCH: an email from a founder describing a startup, seeking investment, or \
introducing a company — especially if it has a PDF pitch deck attached.
NOISE: newsletters, investor updates, cold sales, scheduling, receipts, \
anything clearly not a startup pitch.
UNSURE: ambiguous — use when genuinely uncertain.

Subject: {subject}
From: {sender}
Has PDF: {has_pdf}
Body preview: {snippet}

Reply with one word only: pitch, noise, or unsure"""


def classify_email(raw_email: dict, api_key: str) -> str:
    """Return 'pitch', 'noise', or 'unsure'. Defaults to 'unsure' on any error."""
    has_pdf = any(
        a.get("content_type", "").startswith("application/pdf")
        or a.get("filename", "").lower().endswith(".pdf")
        for a in raw_email.get("attachments", [])
    )
    prompt = _CLASSIFY_PROMPT.format(
        subject=raw_email.get("subject", "")[:200],
        sender=raw_email.get("from", "")[:100],
        has_pdf=str(has_pdf),
        snippet=raw_email.get("body_text", "")[:500],
    )
    resp   = call(HAIKU, [{"role": "user", "content": prompt}], api_key=api_key, max_tokens=10)
    result = text_from(resp).strip().lower().split()[0] if not resp.get("error") else ""
    return result if result in ("pitch", "noise", "unsure") else "unsure"


# ── 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 single JSON object. Be factual and concise. "
    "Cite sources for research findings. Omit fields you cannot determine rather than guessing."
)

_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 produce a JSON evidence packet:
{{
  "company": "Company name",
  "one_liner": "One sentence — what they do",
  "founder_name": "Name or null",
  "founder_email": "email or null",
  "stage": "pre-seed|seed|series_a|series_b|other",
  "sector": "e.g. B2B SaaS / logistics",
  "geography": "City, Country or null",
  "ask": "e.g. $1.5M at $8M cap — or null if not mentioned",
  "traction": "Revenue, users, growth — or 'not mentioned'",
  "team": "Key members and relevant background",
  "business_model": "How they make money",
  "thesis_fit": "Plain English: how well this fits the investor thesis and why",
  "research": [
    {{"finding": "What you found", "source": "URL"}}
  ],
  "missing": ["Important missing information"],
  "questions": ["3-5 sharp questions worth asking before a call"],
  "action": "PASS|REVIEW|ESCALATE",
  "confidence": 0.0,
  "action_reason": "One sentence explaining the triage decision"
}}

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.

confidence is 0.0–1.0.
Return only the JSON object, no markdown, no other text."""


def _triage_with_specialist(packet: dict, thesis: str = "") -> dict:
    """Override action/confidence/action_reason with the bundled Iris specialist.
    Falls back silently to Sonnet's triage if the specialist is not installed."""
    try:
        from specialist.ask import ask  # type: ignore[import]
    except ImportError:
        return packet

    features = {
        "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),
    }
    try:
        result = ask(features, thesis=thesis)
    except Exception:
        return packet

    # Specialist returns natural language — parse action keyword from text
    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


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", []), "pitch.pdf")
    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"}],
        max_tokens=2048, system=_EXTRACT_SYSTEM,
    )
    if resp.get("error"):
        return None
    packet = parse_json(text_from(resp).strip())
    if packet:
        packet = _triage_with_specialist(packet, thesis=thesis)
    return packet


# ── Reply update ───────────────────────────────────────────────────────────────

_UPDATE_SYSTEM = (
    "You are an analyst updating an evidence packet for an early-stage investor. "
    "The investor's screening agent previously sent follow-up questions to the founder. "
    "The founder has now replied. Update the evidence packet with any new information. "
    "Keep all existing fields; refine only what the reply changes or fills in. "
    "Return a single JSON object in the same schema as the original packet."
)

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

Original evidence packet:
{original_packet}

Questions the agent sent to the founder:
{questions}

Founder's reply:
{reply_body}
{attachments_note}
Update the evidence packet with new information from the reply.
Revise action/confidence if warranted.
Remove answered items from missing[], update questions[] to reflect what's still open.
Return only the updated JSON object, no markdown, no other text."""


def update_evidence_from_reply(
    original_packet: dict,
    reply_email: dict,
    thesis: str,
    api_key: str,
) -> dict | None:
    """Update an evidence packet based on the founder's reply, including any attachments."""
    blocks, pdf_count = _attachment_blocks(reply_email.get("attachments", []), "reply.pdf")
    attachments_note  = (
        f"\n{pdf_count} PDF attachment(s) provided above — read them carefully."
        if pdf_count else ""
    )
    questions = original_packet.get("questions") or []
    blocks.append({"type": "text", "text": _UPDATE_PROMPT.format(
        thesis=thesis or "early-stage technology startups",
        original_packet=json.dumps(original_packet, indent=2)[:3000],
        questions="\n".join(f"- {q}" for q in questions),
        reply_body=reply_email.get("body_text", "")[:3000],
        attachments_note=attachments_note,
    )})

    resp = call(
        SONNET, [{"role": "user", "content": blocks}],
        api_key=api_key, max_tokens=2048, system=_UPDATE_SYSTEM,
    )
    if resp.get("error"):
        return None

    updated = parse_json(text_from(resp).strip())
    if updated is None:
        return None

    updated["reply_count"]       = (original_packet.get("reply_count") or 0) + 1
    updated["awaiting_reply"]    = False
    updated["questions_sent_at"] = original_packet.get("questions_sent_at")
    updated = _triage_with_specialist(updated, thesis=thesis)
    return updated


# ── 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)
