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

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

    def record_outcome(self, action: str, result: str, success: bool, strategy: str = '', model: str = '', latency_ms: int = 0, tokens_used: int = 0) -> str:
        content = f"Action: {action}\nResult: {result}"
        return self.vector_store.add(
            content=content,
            metadata={'type': 'outcome', 'success': success, 'strategy': strategy, 'model': model, 'latency_ms': latency_ms, 'tokens_used': tokens_used},
            namespace=self.namespace
        )

    def get_outcomes_by_strategy(self, strategy: str, limit: int = 20) -> list[dict]:
        res = self.vector_store.search_by_metadata({'type': 'outcome', 'strategy': strategy}, namespace=self.namespace, limit=limit)
        return [{'id': r['id'], 'action': r['content'].split('\n')[0].replace('Action: ', ''), 'result': r['content'].split('\n')[1].replace('Result: ', ''), 'success': r['metadata'].get('success'), 'model': r['metadata'].get('model'), 'latency_ms': r['metadata'].get('latency_ms'), 'created_at': r['created_at']} for r in res]

    def get_outcomes_by_model(self, model: str, limit: int = 20) -> list[dict]:
        res = self.vector_store.search_by_metadata({'type': 'outcome', 'model': model}, namespace=self.namespace, limit=limit)
        return [{'id': r['id'], 'action': r['content'].split('\n')[0].replace('Action: ', ''), 'result': r['content'].split('\n')[1].replace('Result: ', ''), 'success': r['metadata'].get('success'), 'strategy': r['metadata'].get('strategy'), 'latency_ms': r['metadata'].get('latency_ms'), 'created_at': r['created_at']} for r in res]

    def success_rate_by_strategy(self) -> dict[str, float]:
        res = self.vector_store.search_by_metadata({'type': 'outcome'}, namespace=self.namespace, limit=10000)
        stats = {}
        for r in res:
            strat = r['metadata'].get('strategy', '')
            if strat not in stats: stats[strat] = {'c': 0, 's': 0}
            stats[strat]['c'] += 1
            if r['metadata'].get('success'): stats[strat]['s'] += 1
        return {k: v['s']/v['c'] for k, v in stats.items()}

    def success_rate_by_model(self) -> dict[str, float]:
        res = self.vector_store.search_by_metadata({'type': 'outcome'}, namespace=self.namespace, limit=10000)
        stats = {}
        for r in res:
            m = r['metadata'].get('model', '')
            if m not in stats: stats[m] = {'c': 0, 's': 0}
            stats[m]['c'] += 1
            if r['metadata'].get('success'): stats[m]['s'] += 1
        return {k: v['s']/v['c'] for k, v in stats.items()}

    def detect_patterns(self, min_occurrences: int = 3) -> list[dict]:
        res = self.vector_store.search_by_metadata({'type': 'outcome'}, namespace=self.namespace, limit=10000)
        stats = {}
        for r in res:
            k = (r['metadata'].get('strategy', ''), r['metadata'].get('model', ''))
            if k not in stats: stats[k] = {'c': 0, 's': 0, 'l': 0}
            stats[k]['c'] += 1
            if r['metadata'].get('success'): stats[k]['s'] += 1
            stats[k]['l'] += r['metadata'].get('latency_ms', 0)
        
        patterns = []
        for (strat, mod), v in stats.items():
            if v['c'] >= min_occurrences:
                patterns.append({'strategy': strat, 'model': mod, 'count': v['c'], 'success_rate': v['s']/v['c'], 'avg_latency_ms': v['l']/v['c']})
        return patterns

    def search_outcomes(self, query: str, success_only: bool = False, top_k: int = 5) -> list[dict]:
        res = self.vector_store.search(query, top_k=top_k*2, namespace=self.namespace)
        filtered = []
        for r in res:
            if success_only and not r['metadata'].get('success'):
                continue
            parts = r['content'].split('\n')
            act = parts[0].replace('Action: ', '') if len(parts) > 0 else ''
            res_str = parts[1].replace('Result: ', '') if len(parts) > 1 else ''
            filtered.append({'id': r['id'], 'action': act, 'result': res_str, 'success': r['metadata'].get('success'), 'score': r['score'], 'created_at': r['created_at']})
            if len(filtered) >= top_k:
                break
        return filtered

    def what_worked_for(self, query: str, top_k: int = 3) -> list[dict]:
        res = self.vector_store.search(query, top_k=top_k*3, namespace=self.namespace)
        out = []
        for r in res:
            if r['metadata'].get('success'):
                parts = r['content'].split('\n')
                out.append({'action': parts[0].replace('Action: ', ''), 'result': parts[1].replace('Result: ', ''), 'strategy': r['metadata'].get('strategy'), 'model': r['metadata'].get('model'), 'score': r['score']})
                if len(out) >= top_k: break
        return out

    def what_failed_for(self, query: str, top_k: int = 3) -> list[dict]:
        res = self.vector_store.search(query, top_k=top_k*3, namespace=self.namespace)
        out = []
        for r in res:
            if not r['metadata'].get('success'):
                parts = r['content'].split('\n')
                out.append({'action': parts[0].replace('Action: ', ''), 'result': parts[1].replace('Result: ', ''), 'strategy': r['metadata'].get('strategy'), 'model': r['metadata'].get('model'), 'score': r['score']})
                if len(out) >= top_k: break
        return out

    def recommend_strategy(self, task_description: str) -> dict | None:
        worked = self.what_worked_for(task_description, top_k=5)
        if not worked:
            return None
        stats = {}
        for w in worked:
            k = (w['strategy'], w['model'])
            stats[k] = stats.get(k, 0) + 1
        best = max(stats.items(), key=lambda x: x[1])
        return {'strategy': best[0][0], 'model': best[0][1], 'confidence': best[1] / len(worked), 'based_on_count': best[1]}

    def context_for_prompt(self, query: str, max_tokens: int = 500) -> str:
        worked = self.what_worked_for(query, top_k=2)
        failed = self.what_failed_for(query, top_k=2)
        max_chars = max_tokens * 4
        context = "[Past outcomes for similar tasks]\n"
        for w in worked:
            add = f"Worked: Action: {w['action']} -> Result: {w['result']}\n"
            if len(context) + len(add) > max_chars: break
            context += add
        for f in failed:
            add = f"Failed: Action: {f['action']} -> Result: {f['result']}\n"
            if len(context) + len(add) > max_chars: break
            context += add
        context += "[End outcome context]\n"
        return context

    def get_stats(self) -> dict:
        c = self.vector_store.conn.cursor()
        c.execute("SELECT metadata FROM vectors WHERE namespace = ?", (self.namespace,))
        tot = 0
        suc = 0
        fail = 0
        strats = set()
        mods = set()
        lat = 0
        for row in c.fetchall():
            tot += 1
            meta = json.loads(row['metadata'] or '{}')
            if meta.get('success'): suc += 1
            else: fail += 1
            if meta.get('strategy'): strats.add(meta['strategy'])
            if meta.get('model'): mods.add(meta['model'])
            lat += meta.get('latency_ms', 0)
        return {
            'total_outcomes': tot, 'success_count': suc, 'failure_count': fail,
            'overall_success_rate': suc / tot if tot > 0 else 0.0,
            'strategies_used': list(strats), 'models_used': list(mods),
            'avg_latency_ms': lat / tot if tot > 0 else 0
        }

    def prune(self, older_than_days: int = 180, keep_failures: bool = True) -> 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'])
            if dt < cutoff:
                meta = json.loads(row['metadata'] or '{}')
                if keep_failures and not meta.get('success'):
                    continue
                to_delete.append(row['id'])
        if to_delete:
            self.vector_store.batch_delete(to_delete)
        return len(to_delete)

    def trend(self, days: int = 30) -> list[dict]:
        cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(days=days)
        c = self.vector_store.conn.cursor()
        c.execute("SELECT metadata, created_at FROM vectors WHERE namespace = ?", (self.namespace,))
        daily = {}
        for row in c.fetchall():
            dt = datetime.datetime.fromisoformat(row['created_at'])
            if dt >= cutoff:
                date_str = dt.strftime('%Y-%m-%d')
                if date_str not in daily: daily[date_str] = {'total': 0, 'successes': 0}
                daily[date_str]['total'] += 1
                meta = json.loads(row['metadata'] or '{}')
                if meta.get('success'): daily[date_str]['successes'] += 1
        trend = []
        for d, v in sorted(daily.items()):
            trend.append({'date': d, 'total': v['total'], 'successes': v['successes'], 'rate': v['successes'] / v['total'] if v['total'] > 0 else 0.0})
        return trend
