from __future__ import annotations

import json
import logging
from typing import Any, Dict, List

import boto3

from storage.config import load_config_from_env
from storage.s3_adapter import S3Storage
from .model import predict

logger = logging.getLogger(__name__)


def lambda_handler(event, context):
    cfg = load_config_from_env()
    s3 = S3Storage(cfg, no_aws=False)
    s3_client = boto3.client("s3", region_name=cfg.region)
    model_version = (event.get("model_version") or "v0.1.0")
    obj = s3_client.get_object(Bucket=cfg.s3_bucket_processed, Key=f"models/{model_version}.joblib")
    model_bytes = obj["Body"].read()
    records: List[Dict[str, Any]] = event.get("records", [])
    texts = [r.get("extracted_text", "") for r in records]
    preds, conf, exp = predict(texts, model_bytes)
    outputs = []
    for r, label, c, weights in zip(records, preds, conf, exp):
        url = r.get("url", "")
        payload = {
            "url": url,
            "classification": label,
            "confidence": c,
            "explain_top_weights": weights,
            "model_version": model_version,
        }
        s3_uri = s3.put_json(cfg.s3_bucket_processed, "classifications", url, payload)
        outputs.append({"s3": s3_uri, "url": url})
    return {"statusCode": 200, "body": json.dumps({"count": len(outputs)})}


