from __future__ import annotations

import io
import json
import logging
import os
from dataclasses import dataclass
from typing import Any, Dict, List, Tuple

import joblib
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline

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

logger = logging.getLogger(__name__)


@dataclass
class ModelArtifact:
    version: str
    s3_uri: str


def build_pipeline() -> Pipeline:
    return Pipeline([
        ("tfidf", TfidfVectorizer(max_features=10000, ngram_range=(1, 2))),
        ("clf", LogisticRegression(max_iter=1000)),
    ])


def train_model(texts: List[str], labels: List[str], version: str = "v0.1.0") -> ModelArtifact:
    pipe = build_pipeline()
    pipe.fit(texts, labels)
    cfg = load_config_from_env()
    s3 = S3Storage(cfg, no_aws=False)
    buffer = io.BytesIO()
    joblib.dump(pipe, buffer)
    payload = buffer.getvalue()
    s3_key_payload = {
        "model_version": version,
        "notes": "scikit-learn logistic regression",
    }
    s3_uri = s3.put_json(cfg.s3_bucket_processed, "models", version, s3_key_payload)
    # store binary alongside manifest
    if not s3.no_aws and s3.s3 is not None:
        obj = s3.s3.Object(cfg.s3_bucket_processed, f"models/{version}.joblib")
        obj.put(Body=payload, ContentType="application/octet-stream", Tagging=f"resource-tag={cfg.resource_tag}")
    return ModelArtifact(version=version, s3_uri=s3_uri)


def predict(texts: List[str], model_bytes: bytes) -> Tuple[List[str], List[float], List[List[Tuple[str, float]]]]:
    pipe: Pipeline = joblib.load(io.BytesIO(model_bytes))
    probs = pipe.predict_proba(texts)
    classes = list(pipe.classes_)
    preds_idx = np.argmax(probs, axis=1)
    preds = [classes[i] for i in preds_idx]
    conf = [float(probs[i, idx]) for i, idx in enumerate(preds_idx)]
    # explainability: use linear weights by feature index
    clf: LogisticRegression = pipe.named_steps["clf"]
    vec: TfidfVectorizer = pipe.named_steps["tfidf"]
    feature_names = np.array(vec.get_feature_names_out())
    explanations: List[List[Tuple[str, float]]] = []
    for i, text in enumerate(texts):
        idx = preds_idx[i]
        # take top positive weights for predicted class
        if clf.coef_.ndim == 2:
            weights = clf.coef_[idx]
        else:
            weights = clf.coef_
        top_indices = np.argsort(weights)[-10:][::-1]
        top_feats = [(feature_names[j].item(), float(weights[j])) for j in top_indices]
        explanations.append(top_feats)
    return preds, conf, explanations


