import json
import datetime
from pathlib import Path
from runtime.memory.vector_store import VectorStore

class ConversationMemory:
    def __init__(self, vector_store: VectorStore, data_dir: Path):
        self.vector_store = vector_store
        self.data_dir = data_dir
        self.namespace = 'conversation'

    def store_summary(self, session_id: str, summary: str, topic: str = '', models_used: list[str] | None = None) -> str:
        return self.vector_store.add(
            content=summary,
            metadata={'type': 'summary', 'session_id': session_id, 'topic': topic, 'models_used': models_used or []},
            namespace=self.namespace
        )

    def extract_decisions(self, conversation_text: str) -> list[str]:
        markers = ['Decision:', 'Decided:', 'We agreed', 'Going with', "Let's use", 'Chose']
        decisions = []
        for line in conversation_text.splitlines():
            line = line.strip()
            for m in markers:
                if line.startswith(m) or line.lower().startswith(m.lower()):
                    decisions.append(line)
                    break
        return decisions

    def extract_action_items(self, conversation_text: str) -> list[str]:
        markers = ['TODO:', 'Action:', 'Next step:', 'Will do:']
        contains = ['need to', 'should', 'must']
        actions = []
        for line in conversation_text.splitlines():
            line = line.strip()
            found = False
            for m in markers:
                if line.startswith(m) or line.lower().startswith(m.lower()):
                    actions.append(line)
                    found = True
                    break
            if not found:
                for c in contains:
                    if c in line.lower():
                        actions.append(line)
                        break
        return actions

    def store_decision(self, session_id: str, decision: str, context: str = '') -> str:
        content = f"Decision: {decision}\nContext: {context}"
        return self.vector_store.add(
            content=content,
            metadata={'type': 'decision', 'session_id': session_id},
            namespace=self.namespace
        )

    def store_action_item(self, session_id: str, action: str, status: str = 'pending') -> str:
        return self.vector_store.add(
            content=action,
            metadata={'type': 'action_item', 'session_id': session_id, 'status': status},
            namespace=self.namespace
        )

    def search_conversations(self, query: str, top_k: int = 5) -> list[dict]:
        res = self.vector_store.search(query, top_k=top_k, namespace=self.namespace)
        return [{'content': r['content'], 'type': r['metadata'].get('type'), 'session_id': r['metadata'].get('session_id'), 'score': r['score'], 'created_at': r['created_at']} for r in res]

    def get_session_memory(self, session_id: str) -> dict:
        res = self.vector_store.search_by_metadata({'session_id': session_id}, namespace=self.namespace, limit=100)
        summary = None
        decisions = []
        action_items = []
        for r in res:
            t = r['metadata'].get('type')
            if t == 'summary':
                summary = r['content']
            elif t == 'decision':
                decisions.append(r)
            elif t == 'action_item':
                action_items.append(r)
        return {'summary': summary, 'decisions': decisions, 'action_items': action_items}

    def get_pending_actions(self) -> list[dict]:
        res = self.vector_store.search_by_metadata({'status': 'pending'}, namespace=self.namespace, limit=100)
        return [{'id': r['id'], 'action': r['content'], 'session_id': r['metadata'].get('session_id'), 'created_at': r['created_at']} for r in res]

    def complete_action(self, action_id: str) -> bool:
        r = self.vector_store.get(action_id)
        if not r:
            return False
        meta = r['metadata']
        meta['status'] = 'completed'
        meta['completed_at'] = datetime.datetime.now(datetime.timezone.utc).isoformat()
        return self.vector_store.update(action_id, metadata=meta)

    def recent_summaries(self, limit: int = 10) -> list[dict]:
        res = self.vector_store.search_by_metadata({'type': 'summary'}, namespace=self.namespace, limit=limit)
        res.sort(key=lambda x: x['created_at'], reverse=True)
        return [{'id': r['id'], 'summary': r['content'], 'topic': r['metadata'].get('topic'), 'session_id': r['metadata'].get('session_id'), 'created_at': r['created_at']} for r in res[:limit]]

    def topic_timeline(self) -> list[dict]:
        res = self.vector_store.search_by_metadata({'type': 'summary'}, namespace=self.namespace, limit=1000)
        topics = {}
        for r in res:
            topic = r['metadata'].get('topic', 'unknown')
            if topic not in topics:
                topics[topic] = {'topic': topic, 'count': 0, 'first_seen': r['created_at'], 'last_seen': r['created_at']}
            topics[topic]['count'] += 1
            topics[topic]['first_seen'] = min(topics[topic]['first_seen'], r['created_at'])
            topics[topic]['last_seen'] = max(topics[topic]['last_seen'], r['created_at'])
        timeline = list(topics.values())
        timeline.sort(key=lambda x: x['last_seen'], reverse=True)
        return timeline

    def apply_decay(self, half_life_days: int = 30) -> int:
        count = 0
        now = datetime.datetime.now(datetime.timezone.utc)
        c = self.vector_store.conn.cursor()
        c.execute('SELECT id, metadata, created_at FROM vectors WHERE namespace = ?', (self.namespace,))
        for row in c.fetchall():
            dt = datetime.datetime.fromisoformat(row['created_at'])
            days_old = (now - dt).days
            if days_old > half_life_days:
                meta = json.loads(row['metadata'] or '{}')
                weight = 0.5 ** (days_old / half_life_days)
                meta['relevance_weight'] = weight
                self.vector_store.update(row['id'], metadata=meta)
                count += 1
        return count

    def context_for_prompt(self, query: str, max_tokens: int = 800) -> str:
        res = self.search_conversations(query, top_k=10)
        max_chars = max_tokens * 4
        context = "[Relevant past conversations]\n"
        for r in res:
            add = f"--- {r['type']} (Session: {r['session_id']}) ---\n{r['content']}\n\n"
            if len(context) + len(add) > max_chars:
                break
            context += add
        context += "[End past context]\n"
        return context

    def summarize_session(self, messages: list[dict]) -> str:
        topic = "General Chat"
        for m in messages:
            if m.get('role') == 'user':
                topic = m.get('content', '')[:50]
                break
        return f"Topic: {topic}. Exchanges: {len(messages)}."

    def auto_memorize(self, session_id: str, messages: list[dict]) -> dict:
        text = "\n".join(m.get('content', '') for m in messages)
        summary = self.summarize_session(messages)
        summary_id = self.store_summary(session_id, summary)
        
        decisions = self.extract_decisions(text)
        decision_ids = [self.store_decision(session_id, d) for d in decisions]
        
        actions = self.extract_action_items(text)
        action_ids = [self.store_action_item(session_id, a) for a in actions]
        
        return {'summary_id': summary_id, 'decision_ids': decision_ids, 'action_ids': action_ids}

    def find_related_sessions(self, session_id: str, top_k: int = 3) -> list[dict]:
        res = self.vector_store.search_by_metadata({'session_id': session_id, 'type': 'summary'}, namespace=self.namespace, limit=1)
        if not res:
            return []
        summary = res[0]['content']
        related = self.vector_store.search(summary, top_k=top_k+1, namespace=self.namespace)
        out = []
        for r in related:
            sid = r['metadata'].get('session_id')
            if sid != session_id and r['metadata'].get('type') == 'summary':
                out.append({'session_id': sid, 'summary': r['content'], 'score': r['score']})
        return out[:top_k]

    def delete_session_memory(self, session_id: str) -> int:
        res = self.vector_store.search_by_metadata({'session_id': session_id}, namespace=self.namespace, limit=1000)
        ids = [r['id'] for r in res]
        if ids:
            self.vector_store.batch_delete(ids)
        return len(ids)

    def get_conversation_stats(self) -> dict:
        c = self.vector_store.conn.cursor()
        c.execute("SELECT metadata FROM vectors WHERE namespace = ?", (self.namespace,))
        sessions = set()
        sums = 0
        decs = 0
        acts = 0
        pend = 0
        topics = set()
        for row in c.fetchall():
            meta = json.loads(row['metadata'] or '{}')
            sid = meta.get('session_id')
            if sid:
                sessions.add(sid)
            t = meta.get('type')
            if t == 'summary':
                sums += 1
                if 'topic' in meta:
                    topics.add(meta['topic'])
            elif t == 'decision':
                decs += 1
            elif t == 'action_item':
                acts += 1
                if meta.get('status') == 'pending':
                    pend += 1
        return {
            'total_sessions': len(sessions),
            'total_summaries': sums,
            'total_decisions': decs,
            'total_actions': acts,
            'pending_actions_count': pend,
            'topics': list(topics)
        }

    def store_insight(self, session_id: str, insight: str, category: str = 'general') -> str:
        return self.vector_store.add(
            content=insight,
            metadata={'type': 'insight', 'session_id': session_id, 'category': category},
            namespace=self.namespace
        )

    def get_insights(self, category: str | None = None, limit: int = 20) -> list[dict]:
        meta = {'type': 'insight'}
        if category:
            meta['category'] = category
        res = self.vector_store.search_by_metadata(meta, namespace=self.namespace, limit=limit)
        return [{'id': r['id'], 'insight': r['content'], 'category': r['metadata'].get('category'), 'session_id': r['metadata'].get('session_id'), 'created_at': r['created_at']} for r in res]

    def export_memory(self) -> dict:
        data = self.vector_store.export_namespace(self.namespace)
        res = {}
        for item in data:
            t = item['metadata'].get('type', 'other')
            if t not in res:
                res[t] = []
            res[t].append(item)
        return res

    def import_memory(self, data: dict) -> int:
        count = 0
        for items in data.values():
            for item in items:
                self.vector_store.add(content=item['content'], metadata=item['metadata'], namespace=self.namespace)
                count += 1
        return count

    def prune(self, older_than_days: int = 180) -> int:
        cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=older_than_days)
        c = self.vector_store.conn.cursor()
        c.execute("SELECT id, metadata, created_at FROM vectors WHERE namespace = ?", (self.namespace,))
        to_delete = []
        for row in c.fetchall():
            dt = datetime.datetime.fromisoformat(row['created_at'])
            meta = json.loads(row['metadata'] or '{}')
            w = meta.get('relevance_weight', 1.0)
            if dt < cutoff or w < 0.1:
                to_delete.append(row['id'])
        if to_delete:
            self.vector_store.batch_delete(to_delete)
        return len(to_delete)

    def search_decisions(self, query: str, top_k: int = 5) -> list[dict]:
        res = self.vector_store.search(query, top_k=top_k * 2, namespace=self.namespace)
        decs = [r for r in res if r['metadata'].get('type') == 'decision']
        return [{'id': r['id'], 'decision': r['content'], 'session_id': r['metadata'].get('session_id'), 'score': r['score'], 'created_at': r['created_at']} for r in decs[:top_k]]
