"""
Iris specialist — local pitch triage using fine-tuned Qwen3-1.7B.

Provides two functions:
  ask(features)          → triage reasoning + decision
  record(features, label) → log feedback for future retraining
"""

from __future__ import annotations

import json
import logging
from pathlib import Path

log = logging.getLogger("superpitch.specialist")

_MODEL = None
_TOKENIZER = None
_MODEL_DIR = Path(__file__).parent

# System prompt must match training data exactly
_SYSTEM_PROMPT = """You are Iris, an AI pitch triage specialist. You evaluate startup pitches on behalf of investors and recommend one of three actions:

- ESCALATE: This deal deserves a meeting. Surface it immediately.
- REVIEW: Worth a closer look, but not urgent. Flag for later review.
- PASS: Not a fit. Skip it.

Your job is to be the investor's first filter. The cost of missing a great deal (false negative) is much higher than taking an extra meeting (false positive). A missed unicorn is catastrophic; a wasted 30-minute call is cheap.

You must consider the investor's specific profile — their stage focus, sector preferences, risk tolerance, and investment style. The SAME deal should get different actions for different investors.

Key principles:
- AI wrappers (thin layers on foundation model APIs) are always PASS — no moat.
- Revenue and capital structure are the strongest triage signals.
- Founder quality (especially previous exits and founder-market fit) is the second tier.
- Different investor types have legitimately different correct answers for the same deal.
- Boring + profitable beats sexy + burning cash.
- High growth in a wrapper doesn't create moat — it attracts competitors.
- Customer concentration is a hidden killer.
- Vertical AI with proprietary data moat is the strongest AI play.
- Previous exit overrides sloppy presentation.
- Moonshots need patient capital — PASS for impatient investors even with great teams."""


def _load_model():
    """Lazy-load the fine-tuned model on first call."""
    global _MODEL, _TOKENIZER

    if _MODEL is not None:
        return

    # Check if model files exist
    model_path = _MODEL_DIR
    if not (model_path / "config.json").exists():
        raise ImportError(
            f"Iris model not found at {model_path}. "
            "Run training/finetune.sh to fine-tune and deploy."
        )

    try:
        from mlx_lm import load
    except ImportError:
        raise ImportError(
            "mlx-lm not installed. Install with: pip install mlx-lm"
        )

    log.info("Loading Iris specialist model from %s", model_path)
    _MODEL, _TOKENIZER = load(str(model_path))
    log.info("Iris specialist loaded")


def _format_prompt(features: dict, thesis: str = "") -> str:
    """Format investor thesis + evidence packet into Iris's expected input format."""
    parts = []

    # Investor context
    if thesis:
        parts.append(f"Investor thesis: {thesis}")
        parts.append("")  # blank line separator

    # Deal description from evidence packet
    if features.get("stage"):
        parts.append(f"Stage: {features['stage']}")
    if features.get("sector"):
        parts.append(f"Sector: {features['sector']}")
    if features.get("team"):
        parts.append(f"Team: {features['team']}")
    if features.get("traction"):
        parts.append(f"Traction: {features['traction']}")
    if features.get("ask"):
        parts.append(f"Ask: {features['ask']}")
    if features.get("thesis_fit"):
        parts.append(f"Thesis fit: {features['thesis_fit']}")
    if features.get("missing"):
        parts.append(f"Missing info: {features['missing']}")
    if features.get("geography"):
        parts.append(f"Geography: {features['geography']}")

    return "\n".join(parts)


def ask(features: dict, thesis: str = "") -> str:
    """
    Triage a pitch using the fine-tuned Iris model.

    Args:
        features: dict with keys from the evidence packet
                  (stage, sector, team, traction, ask, thesis_fit, missing, etc.)
        thesis:   the investor's thesis string (e.g. "early-stage B2B SaaS in fintech")

    Returns:
        str: reasoning text ending with "Decision: ESCALATE/REVIEW/PASS — ..."
    """
    _load_model()

    from mlx_lm import generate

    prompt = _format_prompt(features, thesis=thesis)

    # Build chat messages matching training format
    messages = [
        {"role": "system", "content": _SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]

    # Apply chat template
    formatted = _TOKENIZER.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    response = generate(
        _MODEL, _TOKENIZER,
        prompt=formatted,
        max_tokens=512,
    )

    return response


def record(features: dict, label: str) -> None:
    """
    Log investor decision feedback for future retraining.

    Appends to specialist/feedback.jsonl — each line is:
      {"features": {...}, "investor_decision": "ESCALATE|REVIEW|PASS"}

    This accumulates real decisions that can be used to retrain Iris
    periodically with actual investor preferences.
    """
    feedback_path = _MODEL_DIR / "feedback.jsonl"
    entry = {
        "features": features,
        "investor_decision": label,
    }
    try:
        with open(feedback_path, "a") as f:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except Exception as exc:
        log.warning("Failed to record specialist feedback: %s", exc)
