from __future__ import annotations

import hashlib
import json
import logging
import math
import re
import os
from collections import Counter
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

from config.paths import DATA_ROOT
from receipts.index import search_receipt_vectors
from receipts.models import ReflectionReceipt, RefusalStone
from receipts.writer import write_receipt

REFLECTION_STORE = "reflections.jsonl"
INDEX_FILENAME = "receipt_index.json"
EMBED_DIM = 12
TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
logger = logging.getLogger("loom.receipts.memory")

DEFAULT_VECTOR_INDEX_MAX_AGE_HOURS = 24


def _resolve_vector_index_max_age_hours() -> int:
    token = os.getenv("LOOM_VECTOR_INDEX_MAX_AGE_HOURS")
    if token is None:
        return DEFAULT_VECTOR_INDEX_MAX_AGE_HOURS
    try:
        value = int(str(token))
    except (TypeError, ValueError):
        return DEFAULT_VECTOR_INDEX_MAX_AGE_HOURS
    return max(1, min(value, 168))


VECTOR_INDEX_MAX_AGE_HOURS = _resolve_vector_index_max_age_hours()


def _memory_dir(root: Optional[Path] = None) -> Path:
    base = Path(root) if root is not None else Path(DATA_ROOT)
    target = base / "memory"
    target.mkdir(parents=True, exist_ok=True)
    return target


def _to_iso(value: datetime) -> str:
    if value.tzinfo is None:
        return value.replace(tzinfo=datetime.utcnow().astimezone().tzinfo).isoformat()
    return value.isoformat()


def _from_iso(value: str) -> Optional[datetime]:
    try:
        parsed = datetime.fromisoformat(value)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed.astimezone(timezone.utc)


def _json_default(value: object) -> object:
    if isinstance(value, datetime):
        return _to_iso(value)
    return value

def _scalar(value: Optional[object]) -> Optional[float]:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def tokenize_text(text: str) -> List[str]:
    return [token.lower() for token in TOKEN_RE.findall(text)]


def embed_tokens(tokens: Sequence[str]) -> List[float]:
    if not tokens:
        return [0.0] * EMBED_DIM
    vector = [0.0] * EMBED_DIM
    for token in tokens:
        digest = hashlib.md5(token.encode("utf-8")).hexdigest()
        idx = int(digest[:4], 16) % EMBED_DIM
        sign_bit = int(digest[4:8], 16) & 0x1
        vector[idx] += 1.0 if sign_bit == 0 else -1.0
    norm = math.sqrt(sum(component * component for component in vector))
    if norm == 0.0:
        return vector
    return [component / norm for component in vector]


def cosine_similarity(lhs: Sequence[float], rhs: Sequence[float]) -> float:
    if not lhs or not rhs:
        return 0.0
    length = min(len(lhs), len(rhs))
    dot = sum(lhs[i] * rhs[i] for i in range(length))
    lhs_norm = math.sqrt(sum(lhs[i] * lhs[i] for i in range(length)))
    rhs_norm = math.sqrt(sum(rhs[i] * rhs[i] for i in range(length)))
    if lhs_norm == 0.0 or rhs_norm == 0.0:
        return 0.0
    return dot / (lhs_norm * rhs_norm)


def _write_reflection_store(reflection: ReflectionReceipt, *, root: Optional[Path] = None) -> None:
    directory = _memory_dir(root)
    store_path = directory / REFLECTION_STORE
    payload = reflection.model_dump(by_alias=True)
    with store_path.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(payload, sort_keys=True, default=_json_default))
        handle.write("\n")


def load_recent_reflections(
    domain: str,
    *,
    root: Optional[Path] = None,
    limit: int = 5,
) -> List[ReflectionReceipt]:
    directory = _memory_dir(root)
    store_path = directory / REFLECTION_STORE
    if not store_path.exists():
        return []
    reflections: List[ReflectionReceipt] = []
    with store_path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            try:
                payload = json.loads(line)
            except json.JSONDecodeError:
                continue
            if payload.get("domain") != domain:
                continue
            try:
                reflection = ReflectionReceipt.model_validate(payload)
            except Exception:
                continue
            reflections.append(reflection)
    reflections.sort(key=lambda r: r.window_end, reverse=True)
    return reflections[:limit]


def rollup_reflections(
    refusals: Sequence[RefusalStone],
    *,
    domain: str,
    cohort: str = "global",
    window_start: Optional[datetime] = None,
    window_end: Optional[datetime] = None,
    root: Optional[Path] = None,
) -> ReflectionReceipt:
    if not refusals:
        raise ValueError("At least one refusal receipt is required to build reflections.")

    window_start = window_start or min(refusal.ts for refusal in refusals)
    window_end = window_end or max(refusal.ts for refusal in refusals) + timedelta(seconds=1)

    motif_counter: Counter[str] = Counter()
    ece_values: List[float] = []
    abstain_rates: List[float] = []
    refusal_rates: List[float] = []
    source_ids: List[str] = []

    for refusal in refusals:
        motif_counter.update(refusal.reasons)
        source_ids.append(refusal.receipt_id)
        metrics = refusal.metrics or {}
        ece_val = _scalar(metrics.get("ece_worst") or metrics.get("ece"))
        if ece_val is not None:
            ece_values.append(ece_val)
        abstain_val = _scalar(metrics.get("abstain_rate") or metrics.get("abstention_rate"))
        if abstain_val is not None:
            abstain_rates.append(abstain_val)
        refusal_val = _scalar(metrics.get("refusal_rate"))
        if refusal_val is not None:
            refusal_rates.append(refusal_val)

    top_motifs = [
        {"phrase": phrase, "count": count}
        for phrase, count in motif_counter.most_common(5)
    ]

    metrics_summary: Dict[str, Optional[float]] = {
        "ece_avg": float(sum(ece_values) / len(ece_values)) if ece_values else None,
        "abstain_rate_avg": float(sum(abstain_rates) / len(abstain_rates)) if abstain_rates else None,
        "refusal_rate_avg": float(sum(refusal_rates) / len(refusal_rates)) if refusal_rates else None,
    }

    if ece_values:
        metrics_summary["ece_trend"] = float(ece_values[-1] - ece_values[0])

    summary_parts: List[str] = []
    if top_motifs:
        primary = top_motifs[0]
        summary_parts.append(f"Top motif: '{primary['phrase']}' ({primary['count']}x).")
    if metrics_summary.get("ece_trend") is not None:
        trend = metrics_summary["ece_trend"]
        direction = "up" if trend and trend > 0 else "down"
        summary_parts.append(f"ECE trend {direction} {abs(trend):.3f} over window.")
    if metrics_summary.get("abstain_rate_avg") is not None:
        summary_parts.append(f"Abstain rate avg {metrics_summary['abstain_rate_avg']:.3f}.")
    reflection_summary = " ".join(summary_parts) if summary_parts else "No significant changes detected."

    reflection = ReflectionReceipt(
        receipt_id=f"REFLECTION_{domain.upper()}_{window_start.strftime('%Y%m%d')}",
        domain=domain,
        cohort=cohort,
        window_start=window_start,
        window_end=window_end,
        refusal_motifs=top_motifs,
        metrics_summary={key: value for key, value in metrics_summary.items() if value is not None},
        reflection_summary=reflection_summary,
        source_receipt_ids=source_ids,
    )
    write_receipt(reflection, root=root)
    _write_reflection_store(reflection, root=root)
    return reflection


def _load_receipt_index(root: Optional[Path] = None) -> Dict[str, object]:
    directory = _memory_dir(root)
    index_path = directory / INDEX_FILENAME
    if not index_path.exists():
        return {}
    try:
        payload = json.loads(index_path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}
    return payload if isinstance(payload, dict) else {}


def generate_refusal_hint_payload(
    *,
    domain: str,
    reasons: Sequence[str],
    metrics: Optional[Mapping[str, object]] = None,
    asset_id: Optional[str] = None,
    cohort: Optional[str] = None,
    root: Optional[Path] = None,
    top_k: int = 3,
) -> Dict[str, object]:
    reasons_text = " ".join(reasons)
    tokens = tokenize_text(reasons_text)
    query_vector = embed_tokens(tokens)
    index_payload = _load_receipt_index(root=root)
    items = index_payload.get("items") if isinstance(index_payload, dict) else None
    alerts: List[str] = []
    memory_alerts: List[str] = []
    def _record_alert(message: str) -> None:
        alerts.append(message)
        memory_alerts.append(message)
    vector_index_generated_at: Optional[str] = None
    vector_index_age_hours: Optional[float] = None
    policy_hours = VECTOR_INDEX_MAX_AGE_HOURS
    similar_refs: List[Dict[str, object]] = []

    if index_payload:
        raw_generated_at = index_payload.get("generated_at")
        generated_dt: Optional[datetime] = None
        if isinstance(raw_generated_at, str):
            parsed_dt = _from_iso(raw_generated_at)
            if parsed_dt is not None:
                generated_dt = parsed_dt
                vector_index_generated_at = parsed_dt.isoformat()
            else:
                _record_alert("Vector index timestamp malformed; treating as stale.")
        elif raw_generated_at is not None:
            _record_alert("Vector index timestamp missing; treating as stale.")

        if generated_dt is not None:
            now = datetime.now(timezone.utc)
            age_hours = max(0.0, (now - generated_dt).total_seconds() / 3600.0)
            vector_index_age_hours = age_hours
            if policy_hours and age_hours > policy_hours:
                _record_alert(
                    f"Vector index stale ({age_hours:.1f}h > policy {policy_hours}h); nightly refresh required."
                )
        elif not alerts:
            _record_alert("Vector index age unknown; nightly refresh status unavailable.")
    else:
        _record_alert("Vector index missing; semantic recall may be degraded until refresh completes.")

    vector_hits: List[Dict[str, Any]] = []
    try:
        vector_hits = search_receipt_vectors(
            query_vector,
            root=root,
            domain=domain,
            receipt_types=("RefusalStone",),
            exclude_asset_ids=[asset_id] if asset_id else None,
            limit=max(top_k or 0, 6),
            min_score=0.05,
        )
    except Exception:  # pragma: no cover - defensive logging path
        logger.exception("semantic receipt search failed", extra={"domain": domain})
        vector_hits = []

    for match in vector_hits:
        similar_refs.append(
            {
                "receipt_id": match.get("receipt_id"),
                "summary": match.get("summary"),
                "score": round(float(match.get("score", 0.0)), 4),
                "type": match.get("type"),
            }
        )

    use_legacy_fallback = (
        isinstance(items, list)
        and (
            not similar_refs
            or (top_k > 0 and len(similar_refs) < top_k)
        )
    )

    if use_legacy_fallback:
        for item in items:
            symbolic = item.get("symbolic", {})
            if not isinstance(symbolic, dict):
                continue
            if symbolic.get("domain") != domain:
                continue
            if asset_id and symbolic.get("asset_id") == asset_id:
                continue
            vector = item.get("vector")
            if not isinstance(vector, list):
                continue
            score = cosine_similarity(query_vector, vector)
            if score <= 0:
                continue
            similar_refs.append(
                {
                    "receipt_id": symbolic.get("receipt_id"),
                    "summary": item.get("summary"),
                    "score": round(float(score), 4),
                    "type": symbolic.get("type"),
                }
            )

    reflections = load_recent_reflections(domain, root=root, limit=5)
    reflection_hint: Optional[str] = None
    reflection_receipt_id: Optional[str] = None
    if reflections:
        query_terms = set(tokens)
        for reflection in reflections:
            motifs = reflection.refusal_motifs or []
            motif_terms: set[str] = set()
            for motif in motifs:
                if isinstance(motif, Mapping):
                    phrase = str(motif.get("phrase") or "")
                    motif_terms.update(tokenize_text(phrase))
            if query_terms.intersection(motif_terms):
                reflection_hint = reflection.reflection_summary
                reflection_receipt_id = reflection.receipt_id
                similar_refs.insert(
                    0,
                    {
                        "receipt_id": reflection.receipt_id,
                        "summary": reflection.reflection_summary,
                        "score": 0.75,
                        "type": "ReflectionReceipt",
                    },
                )
                break

    # Deduplicate similar refs by receipt_id preserving highest score.
    unique_refs: Dict[str, Dict[str, object]] = {}
    for entry in similar_refs:
        receipt_id = entry.get("receipt_id")
        if not receipt_id:
            continue
        existing = unique_refs.get(receipt_id)
        if existing is None or entry.get("score", 0.0) > existing.get("score", 0.0):
            unique_refs[receipt_id] = entry

    ordered_refs = list(unique_refs.values())
    ordered_refs.sort(key=lambda entry: entry.get("score", 0.0), reverse=True)
    if top_k > 0:
        ordered_refs = ordered_refs[:top_k]
    if reflection_receipt_id:
        reflection_entry = unique_refs.get(reflection_receipt_id)
        if reflection_entry and reflection_entry not in ordered_refs:
            ordered_refs = [reflection_entry] + ordered_refs
            if top_k > 0:
                ordered_refs = ordered_refs[:top_k]

    payload: Dict[str, Any] = {
        "similar_case_refs": ordered_refs,
        "reflection_hint": reflection_hint,
        "alerts": alerts,
        "memory_alerts": memory_alerts,
    }
    if vector_index_generated_at is not None:
        payload["vector_index_generated_at"] = vector_index_generated_at
    if vector_index_age_hours is not None:
        payload["vector_index_age_hours"] = vector_index_age_hours
    if policy_hours:
        payload["vector_index_policy_hours"] = policy_hours
    return payload
