"""
Memory store for CrashSense — persists every diagnosis, fix, and AI action
in a local SQLite database.

Tables
------
- ``diagnoses``   – one row per analysed transaction / raw-error session.
- ``ai_actions``  – granular log of every LLM call, pattern match, and
                    remediation step (foreign-keyed to diagnoses).

Public helpers convert rows to plain dicts so the REST API (or any downstream
code) can serialise them without importing SQLAlchemy models.
"""

from __future__ import annotations

import hashlib
import json
import os
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional

from sqlalchemy import (
    Column, DateTime, Float, ForeignKey, Integer, String, Text,
    create_engine, func,
)
from sqlalchemy.orm import declarative_base, relationship, sessionmaker

from ..config import load_config

Base = declarative_base()


# ── ORM Models ────────────────────────────────────────────────────────────

class DiagnosisRecord(Base):
    """Full diagnosis record — one per tx_hash or raw-error session."""

    __tablename__ = "diagnoses"

    id              = Column(Integer, primary_key=True)
    tx_hash         = Column(String(66), index=True, unique=True)
    chain           = Column(String(32), default="ethereum")
    problem_hash    = Column(String(64), index=True)

    # ── AI analysis ───────────────────────────────────────────────────────
    summary         = Column(Text, default="")
    root_cause      = Column(Text, default="")
    suggested_fix   = Column(Text, default="")
    explanation     = Column(Text, default="")

    # ── Patterns & remediation ────────────────────────────────────────────
    patterns_json   = Column(Text, default="[]")   # JSON list of pattern strings
    remediation_json = Column(Text, default="[]")  # JSON list of remediation dicts

    # ── LLM metadata ─────────────────────────────────────────────────────
    provider        = Column(String(32), default="")
    model           = Column(String(64), default="")

    # ── Transaction context (denormalised for fast queries) ───────────────
    status          = Column(String(16), default="")
    function_called = Column(String(128), default="")
    from_address    = Column(String(42), default="")
    to_address      = Column(String(42), default="")
    value_eth       = Column(Float, default=0.0)
    gas_used        = Column(Integer, default=0)
    gas_limit       = Column(Integer, default=0)
    revert_reason   = Column(Text, default="")

    # ── Full JSON blobs for anything we don't denormalise ─────────────────
    tx_context_json = Column(Text, default="{}")   # raw ChainFetcher output
    analysis_json   = Column(Text, default="{}")   # raw LLM output
    raw_input       = Column(Text, default="")      # for diagnose-raw

    # ── Bookkeeping ──────────────────────────────────────────────────────
    frequency       = Column(Integer, default=1)
    created_at      = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    last_accessed   = Column(DateTime, default=lambda: datetime.now(timezone.utc))

    # ── Relationships ────────────────────────────────────────────────────
    actions = relationship("AIAction", back_populates="diagnosis", cascade="all, delete-orphan")


class AIAction(Base):
    """Granular log entry for every AI/engine action during a diagnosis."""

    __tablename__ = "ai_actions"

    id            = Column(Integer, primary_key=True)
    diagnosis_id  = Column(Integer, ForeignKey("diagnoses.id"), index=True)
    action_type   = Column(String(32))   # "pattern_match" | "llm_call" | "remediation" | "fetch"
    detail        = Column(Text, default="")   # human-readable description
    data_json     = Column(Text, default="{}")  # structured payload
    created_at    = Column(DateTime, default=lambda: datetime.now(timezone.utc))

    diagnosis = relationship("DiagnosisRecord", back_populates="actions")


# ── Backward-compat alias used by tests ──────────────────────────────────
DiagnosisMemory = DiagnosisRecord


# ── MemoryStore ───────────────────────────────────────────────────────────

class MemoryStore:
    """SQLite-backed store for all CrashSense history."""

    def __init__(self, db_path: str | None = None) -> None:
        if db_path is None:
            cfg = load_config()
            db_path = cfg.get("memory", {}).get("path", "")
            if not db_path:
                from pathlib import Path
                db_path = str(Path.home() / ".crashsense" / "memories.db")
        parent = os.path.dirname(db_path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        self.engine = create_engine(
            f"sqlite:///{db_path}",
            connect_args={"check_same_thread": False},
        )
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)

    # ── Save (upsert) ────────────────────────────────────────────────────

    def save_diagnosis(
        self,
        tx_hash: str,
        chain: str = "ethereum",
        summary: str = "",
        root_cause: str = "",
        suggested_fix: str = "",
        status: str = "",
        function_called: str = "",
        *,
        explanation: str = "",
        patterns: list | None = None,
        remediation: list | None = None,
        provider: str = "",
        model: str = "",
        from_address: str = "",
        to_address: str = "",
        value_eth: float = 0.0,
        gas_used: int = 0,
        gas_limit: int = 0,
        revert_reason: str = "",
        tx_context: dict | None = None,
        analysis: dict | None = None,
        raw_input: str = "",
    ) -> DiagnosisRecord:
        """Insert or update (upsert) a full diagnosis record and return it."""
        h = self._hash(tx_hash)
        session = self.Session()
        try:
            rec = session.query(DiagnosisRecord).filter_by(tx_hash=tx_hash).one_or_none()
            if rec:
                rec.frequency += 1
                rec.last_accessed = datetime.now(timezone.utc)
                rec.summary = summary or rec.summary
                rec.root_cause = root_cause or rec.root_cause
                rec.suggested_fix = suggested_fix or rec.suggested_fix
                rec.explanation = explanation or rec.explanation
                if patterns is not None:
                    rec.patterns_json = json.dumps(patterns)
                if remediation is not None:
                    rec.remediation_json = json.dumps(remediation)
                if tx_context is not None:
                    rec.tx_context_json = json.dumps(tx_context, default=str)
                if analysis is not None:
                    rec.analysis_json = json.dumps(analysis, default=str)
                if provider:
                    rec.provider = provider
                if model:
                    rec.model = model
            else:
                rec = DiagnosisRecord(
                    tx_hash=tx_hash, chain=chain, problem_hash=h,
                    summary=summary, root_cause=root_cause,
                    suggested_fix=suggested_fix, explanation=explanation,
                    patterns_json=json.dumps(patterns or []),
                    remediation_json=json.dumps(remediation or []),
                    provider=provider, model=model,
                    status=status, function_called=function_called,
                    from_address=from_address, to_address=to_address,
                    value_eth=value_eth, gas_used=gas_used, gas_limit=gas_limit,
                    revert_reason=revert_reason,
                    tx_context_json=json.dumps(tx_context or {}, default=str),
                    analysis_json=json.dumps(analysis or {}, default=str),
                    raw_input=raw_input,
                    last_accessed=datetime.now(timezone.utc),
                )
                session.add(rec)
            session.commit()
            session.refresh(rec)
            return rec
        finally:
            session.close()

    # ── Log an AI action ─────────────────────────────────────────────────

    def log_action(
        self,
        diagnosis_id: int,
        action_type: str,
        detail: str = "",
        data: dict | None = None,
    ) -> AIAction:
        """Append a granular action log entry to a diagnosis."""
        session = self.Session()
        try:
            action = AIAction(
                diagnosis_id=diagnosis_id,
                action_type=action_type,
                detail=detail,
                data_json=json.dumps(data or {}, default=str),
            )
            session.add(action)
            session.commit()
            session.refresh(action)
            return action
        finally:
            session.close()

    # ── Read ─────────────────────────────────────────────────────────────

    def get_by_tx_hash(self, tx_hash: str) -> Optional[DiagnosisRecord]:
        session = self.Session()
        try:
            return session.query(DiagnosisRecord).filter_by(tx_hash=tx_hash).one_or_none()
        finally:
            session.close()

    def list_recent(self, limit: int = 50) -> List[DiagnosisRecord]:
        session = self.Session()
        try:
            return (
                session.query(DiagnosisRecord)
                .order_by(DiagnosisRecord.last_accessed.desc())
                .limit(limit)
                .all()
            )
        finally:
            session.close()

    def search(self, keyword: str, limit: int = 20) -> List[DiagnosisRecord]:
        session = self.Session()
        try:
            like = f"%{keyword}%"
            return (
                session.query(DiagnosisRecord)
                .filter(
                    (DiagnosisRecord.summary.ilike(like))
                    | (DiagnosisRecord.root_cause.ilike(like))
                    | (DiagnosisRecord.tx_hash.ilike(like))
                )
                .order_by(DiagnosisRecord.last_accessed.desc())
                .limit(limit)
                .all()
            )
        finally:
            session.close()

    def get_all(self, page: int = 1, per_page: int = 50) -> Dict[str, Any]:
        """Paginated fetch of all diagnoses (newest first)."""
        session = self.Session()
        try:
            total = session.query(func.count(DiagnosisRecord.id)).scalar() or 0
            items = (
                session.query(DiagnosisRecord)
                .order_by(DiagnosisRecord.last_accessed.desc())
                .offset((page - 1) * per_page)
                .limit(per_page)
                .all()
            )
            return {
                "total": total,
                "page": page,
                "per_page": per_page,
                "pages": max(1, (total + per_page - 1) // per_page),
                "items": items,
            }
        finally:
            session.close()

    def get_actions(self, diagnosis_id: int) -> List[AIAction]:
        """Return all AI actions for a diagnosis."""
        session = self.Session()
        try:
            return (
                session.query(AIAction)
                .filter_by(diagnosis_id=diagnosis_id)
                .order_by(AIAction.created_at)
                .all()
            )
        finally:
            session.close()

    def get_stats(self) -> Dict[str, Any]:
        """Aggregate statistics across all diagnoses."""
        session = self.Session()
        try:
            total = session.query(func.count(DiagnosisRecord.id)).scalar() or 0
            chains = session.query(DiagnosisRecord.chain, func.count(DiagnosisRecord.id)).group_by(DiagnosisRecord.chain).all()
            providers = session.query(DiagnosisRecord.provider, func.count(DiagnosisRecord.id)).group_by(DiagnosisRecord.provider).all()
            statuses = session.query(DiagnosisRecord.status, func.count(DiagnosisRecord.id)).group_by(DiagnosisRecord.status).all()
            return {
                "total_diagnoses": total,
                "by_chain": {c: n for c, n in chains if c},
                "by_provider": {p: n for p, n in providers if p},
                "by_status": {s: n for s, n in statuses if s},
            }
        finally:
            session.close()

    def prune(self, max_entries: int = 2000, retention_days: int = 365) -> int:
        session = self.Session()
        try:
            total = session.query(func.count(DiagnosisRecord.id)).scalar() or 0
            deleted = 0
            if total <= max_entries:
                return 0
            cutoff = datetime.now(timezone.utc) - timedelta(days=retention_days)
            deleted = session.query(DiagnosisRecord).filter(DiagnosisRecord.last_accessed < cutoff).delete()
            total = session.query(func.count(DiagnosisRecord.id)).scalar() or 0
            if total > max_entries:
                excess = total - max_entries
                oldest = (
                    session.query(DiagnosisRecord)
                    .order_by(DiagnosisRecord.last_accessed)
                    .limit(excess)
                    .all()
                )
                for o in oldest:
                    session.delete(o)
                    deleted += 1
            session.commit()
            return deleted
        finally:
            session.close()

    # ── Serialisation helpers (for REST API / SDK) ───────────────────────

    @staticmethod
    def record_to_dict(rec: DiagnosisRecord) -> Dict[str, Any]:
        """Convert a DiagnosisRecord to a JSON-safe dict."""
        return {
            "id": rec.id,
            "tx_hash": rec.tx_hash,
            "chain": rec.chain,
            "summary": rec.summary,
            "root_cause": rec.root_cause,
            "suggested_fix": rec.suggested_fix,
            "explanation": rec.explanation,
            "patterns": json.loads(rec.patterns_json or "[]"),
            "remediation": json.loads(rec.remediation_json or "[]"),
            "provider": rec.provider,
            "model": rec.model,
            "status": rec.status,
            "function_called": rec.function_called,
            "from_address": rec.from_address,
            "to_address": rec.to_address,
            "value_eth": rec.value_eth,
            "gas_used": rec.gas_used,
            "gas_limit": rec.gas_limit,
            "revert_reason": rec.revert_reason,
            "tx_context": json.loads(rec.tx_context_json or "{}"),
            "analysis": json.loads(rec.analysis_json or "{}"),
            "raw_input": rec.raw_input,
            "frequency": rec.frequency,
            "created_at": rec.created_at.isoformat() if rec.created_at else None,
            "last_accessed": rec.last_accessed.isoformat() if rec.last_accessed else None,
        }

    @staticmethod
    def action_to_dict(action: AIAction) -> Dict[str, Any]:
        return {
            "id": action.id,
            "diagnosis_id": action.diagnosis_id,
            "action_type": action.action_type,
            "detail": action.detail,
            "data": json.loads(action.data_json or "{}"),
            "created_at": action.created_at.isoformat() if action.created_at else None,
        }

    # ── Internal ─────────────────────────────────────────────────────────

    @staticmethod
    def _hash(text: str) -> str:
        return hashlib.sha256(text.encode("utf-8")).hexdigest()
