"""Shopper prediction via BYOK HF/OpenAI APIs, vision or text."""

from __future__ import annotations

import argparse
import base64
import json
import os
import sys
from pathlib import Path
from typing import Any, Literal

import pandas as pd

from sicss2026.html_text import html_to_text
from sicss2026.kit_paths import (
    default_persona_prompts,
    default_scenarios_dir,
    default_screenshots_dir,
)
from sicss2026.prediction_parser import parse_prediction

DEFAULT_VL_MODEL = "Qwen/Qwen3-VL-30B-A3B-Instruct"
DEFAULT_TEXT_MODEL = "Qwen/Qwen3-30B-A3B-Instruct"
DEFAULT_HF_ROUTER_URL = "https://router.huggingface.co/v1"

OUTPUT_FORMAT = """Respond with JSON only:
{
  "record_id": "<id>",
  "persona_id": "<persona_id>",
  "purchase_probability": <number between 0 and 1>,
  "reason": "<one sentence grounded in visible page content>"
}"""


def resolve_api_config() -> tuple[str, str, Literal["byok"]]:
    api_key = (
        os.environ.get("OPENAI_API_KEY")
        or os.environ.get("HF_TOKEN")
        or ""
    ).strip()
    if not api_key:
        raise RuntimeError(
            "Set HF_TOKEN (Hugging Face) or OPENAI_API_KEY (OpenAI) to run inference."
        )
    base_url = (
        os.environ.get("INFERENCE_BASE_URL")
        or os.environ.get("OPENAI_BASE_URL")
        or DEFAULT_HF_ROUTER_URL
    ).rstrip("/")
    return base_url, api_key, "byok"


def default_model_for_mode(mode: str) -> str:
    if mode == "text":
        return DEFAULT_TEXT_MODEL
    return DEFAULT_VL_MODEL


def load_persona_prompt(prompts_dir: Path, persona_id: str) -> str:
    path = prompts_dir / f"{persona_id}.txt"
    if not path.exists():
        raise FileNotFoundError(f"Missing persona prompt: {path}")
    return path.read_text(encoding="utf-8").strip()


def resolve_screenshot(screenshots_dir: Path, record_id: str) -> Path:
    path = screenshots_dir / f"{record_id}.png"
    if not path.exists():
        raise FileNotFoundError(f"Screenshot not found: {path}")
    return path


def resolve_scenario_html(scenarios_dir: Path, record_id: str) -> Path:
    path = scenarios_dir / f"{record_id}.html"
    if not path.exists():
        raise FileNotFoundError(f"Scenario HTML not found: {path}")
    return path


def encode_image(path: Path) -> str:
    data = path.read_bytes()
    b64 = base64.standard_b64encode(data).decode("ascii")
    return f"data:image/png;base64,{b64}"


def build_vision_messages(
    *,
    persona_prompt: str,
    persona_id: str,
    persona_name: str,
    record_id: str,
    image_data_url: str,
) -> list[dict[str, Any]]:
    system = (
        "You are a synthetic shopper agent. Inspect the product page screenshot and "
        "estimate purchase probability. Use only visible information. /no_think"
    )
    user_text = (
        f"Shopper persona:\n{persona_prompt}\n\n"
        f"Persona ID: {persona_id}\n"
        f"Persona name: {persona_name}\n"
        f"Record ID: {record_id}\n\n"
        f"{OUTPUT_FORMAT}"
    )
    return [
        {"role": "system", "content": system},
        {
            "role": "user",
            "content": [
                {"type": "text", "text": user_text},
                {"type": "image_url", "image_url": {"url": image_data_url}},
            ],
        },
    ]


def build_text_messages(
    *,
    persona_prompt: str,
    persona_id: str,
    persona_name: str,
    record_id: str,
    page_text: str,
) -> list[dict[str, Any]]:
    system = (
        "You are a synthetic shopper agent. Read the product page text and estimate "
        "purchase probability. Use only information present in the page text. /no_think"
    )
    user_text = (
        f"Shopper persona:\n{persona_prompt}\n\n"
        f"Persona ID: {persona_id}\n"
        f"Persona name: {persona_name}\n"
        f"Record ID: {record_id}\n\n"
        f"Product page text:\n{page_text}\n\n"
        f"{OUTPUT_FORMAT}"
    )
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": user_text},
    ]


def chat_completion(
    *,
    messages: list[dict[str, Any]],
    model: str,
    max_tokens: int = 256,
) -> dict[str, Any]:
    try:
        from openai import OpenAI
    except ImportError as exc:
        raise RuntimeError("Install openai: poetry install") from exc

    base_url, api_key, auth_mode = resolve_api_config()
    client = OpenAI(base_url=base_url, api_key=api_key)
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        max_tokens=max_tokens,
        temperature=0.3,
    )
    choice = response.choices[0]
    usage = response.usage
    return {
        "content": choice.message.content or "",
        "usage": usage.model_dump() if usage else {},
        "model": response.model,
        "auth_mode": auth_mode,
        "base_url": base_url,
    }


def predict_record(
    *,
    record_id: str,
    persona_id: str,
    persona_name: str,
    persona_prompt: str,
    mode: str,
    screenshots_dir: Path,
    scenarios_dir: Path,
    model: str,
) -> dict[str, Any]:
    if mode == "text":
        page_text = html_to_text(resolve_scenario_html(scenarios_dir, record_id))
        messages = build_text_messages(
            persona_prompt=persona_prompt,
            persona_id=persona_id,
            persona_name=persona_name,
            record_id=record_id,
            page_text=page_text,
        )
    else:
        messages = build_vision_messages(
            persona_prompt=persona_prompt,
            persona_id=persona_id,
            persona_name=persona_name,
            record_id=record_id,
            image_data_url=encode_image(resolve_screenshot(screenshots_dir, record_id)),
        )
    result = chat_completion(messages=messages, model=model)
    parsed = parse_prediction(result["content"])
    return {
        "record_id": record_id,
        "persona_id": persona_id,
        **parsed,
        "model": result.get("model"),
        "mode": mode,
        "usage": result.get("usage"),
        "auth_mode": result.get("auth_mode"),
    }


def cmd_predict(args: argparse.Namespace) -> int:
    prompts_dir = args.persona_prompts
    persona_prompt = (
        args.persona_prompt.read_text(encoding="utf-8").strip()
        if args.persona_prompt
        else load_persona_prompt(prompts_dir, args.persona_id)
    )
    model = args.model or default_model_for_mode(args.mode)
    prediction = predict_record(
        record_id=args.record_id,
        persona_id=args.persona_id,
        persona_name=args.persona_name,
        persona_prompt=persona_prompt,
        mode=args.mode,
        screenshots_dir=args.screenshots_dir,
        scenarios_dir=args.scenarios_dir,
        model=model,
    )
    print(json.dumps(prediction, indent=2))
    if args.trace:
        args.trace.parent.mkdir(parents=True, exist_ok=True)
        args.trace.write_text(json.dumps(prediction, indent=2), encoding="utf-8")
    return 0


def cmd_batch(args: argparse.Namespace) -> int:
    df = pd.read_csv(args.input)
    required = {"record_id", "persona_id", "persona_name"}
    missing = required - set(df.columns)
    if missing:
        raise SystemExit(f"CSV missing columns: {sorted(missing)}")

    model = args.model or default_model_for_mode(args.mode)
    outputs: list[dict[str, Any]] = []
    for row in df.itertuples(index=False):
        persona_prompt = load_persona_prompt(args.persona_prompts, row.persona_id)
        prediction = predict_record(
            record_id=str(row.record_id),
            persona_id=str(row.persona_id),
            persona_name=str(row.persona_name),
            persona_prompt=persona_prompt,
            mode=args.mode,
            screenshots_dir=args.screenshots_dir,
            scenarios_dir=args.scenarios_dir,
            model=model,
        )
        outputs.append(prediction)
        print(json.dumps(prediction), flush=True)

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", encoding="utf-8") as handle:
            for item in outputs:
                handle.write(json.dumps(item) + "\n")
        print(f"Wrote {len(outputs)} predictions to {args.output}", file=sys.stderr)
    return 0


def add_shared_options(parser: argparse.ArgumentParser, *, as_override: bool = False) -> None:
    """Register options shared by the top-level parser and each subcommand.

    When ``as_override`` is True the options use ``argparse.SUPPRESS`` defaults so
    that omitting them after the subcommand leaves the value parsed before it
    untouched. This lets ``--mode`` (etc.) be given on either side of the
    subcommand, e.g. both ``shop-vl-predict --mode text batch ...`` and
    ``shop-vl-predict batch --mode text ...`` work.
    """
    parser.add_argument(
        "--mode",
        choices=("vision", "text"),
        default=argparse.SUPPRESS if as_override else "vision",
        help="vision: screenshot + VL model; text: HTML page text + text model",
    )
    parser.add_argument(
        "--model",
        default=argparse.SUPPRESS if as_override else None,
        help="Model id (defaults: Qwen3-VL for vision, Qwen3-30B for text)",
    )
    parser.add_argument(
        "--screenshots-dir",
        type=Path,
        default=argparse.SUPPRESS if as_override else default_screenshots_dir(),
    )
    parser.add_argument(
        "--scenarios-dir",
        type=Path,
        default=argparse.SUPPRESS if as_override else default_scenarios_dir(),
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Shopper prediction via your own API key (BYOK)."
    )
    add_shared_options(parser)
    sub = parser.add_subparsers(dest="command", required=True)

    single = sub.add_parser("predict", help="Predict one scenario")
    add_shared_options(single, as_override=True)
    single.add_argument("--record-id", required=True)
    single.add_argument("--persona-id", required=True)
    single.add_argument("--persona-name", required=True)
    single.add_argument("--persona-prompts", type=Path, default=default_persona_prompts())
    single.add_argument("--persona-prompt", type=Path, help="Override persona prompt file")
    single.add_argument("--trace", type=Path)
    single.set_defaults(func=cmd_predict)

    batch = sub.add_parser("batch", help="Predict all rows in a CSV")
    add_shared_options(batch, as_override=True)
    batch.add_argument("--input", type=Path, required=True)
    batch.add_argument("--persona-prompts", type=Path, default=default_persona_prompts())
    batch.add_argument("--output", type=Path)
    batch.set_defaults(func=cmd_batch)
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    try:
        return args.func(args)
    except (RuntimeError, FileNotFoundError, ValueError) as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
