#!/usr/bin/env python3
"""
One-shot idempotent migration script for Phase 1: Dual-Write (JSONL + SQLite).
"""
import sys
import argparse
import json
from pathlib import Path

from pulse.core.brain_db import get_conn, init_db, insert_knowledge, DB_PATH
from pulse.core.brain_utils import (
    LESSONS_DIR, FAILS_DIR, PATTERNS_DIR, PRIVATE_DIR, INTENTS_DIR, DECISIONS_DIR,
    EVIDENCE_DIR, KNOWLEDGE_DIR, GRAPH_FILE
)

def _rollback():
    """Drops the entire SQLite file."""
    conn = get_conn()
    conn.close()
    if DB_PATH.exists():
        DB_PATH.unlink()
        print(f"Rollback complete: {DB_PATH} deleted.")
    else:
        print("No database found to rollback.")

def migrate_jsonl(table: str, path: Path, k_type: str) -> int:
    """Migrate a standard JSONL file to a table."""
    if not path.exists():
        return 0
    count = 0
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
                data["type"] = k_type
                insert_knowledge(table, data)
                count += 1
            except Exception as e:
                print(f"Error migrating {path.name}: {e}")
    return count

def migrate_json(table: str, path: Path, k_type: str) -> int:
    """Migrate a standard JSON list-of-dicts to a table."""
    if not path.exists():
        return 0
    count = 0
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                for item in data:
                    item["type"] = k_type
                    if "content" not in item and "fact" in item:
                        item["content"] = item["fact"]
                    insert_knowledge(table, item)
                    count += 1
    except Exception as e:
        print(f"Error migrating {path.name}: {e}")
    return count

def migrate_retrieval_metrics(path: Path) -> int:
    """Migrates retrieval metrics JSON map."""
    if not path.exists():
        return 0
    count = 0
    conn = get_conn()
    conn.execute("CREATE TABLE IF NOT EXISTS retrieval_metrics (item_key TEXT PRIMARY KEY, retrieval_count INTEGER DEFAULT 0, last_accessed TEXT);")
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            with conn:
                for k, v in data.get("items", {}).items():
                    access_count = v if isinstance(v, int) else v.get("access_count", 0)
                    last_accessed = None if isinstance(v, int) else v.get("last_accessed")
                    conn.execute("INSERT OR IGNORE INTO retrieval_metrics (item_key, retrieval_count, last_accessed) VALUES (?, ?, ?)", 
                                 (k, access_count, last_accessed))
                    count += 1
    except Exception as e:
        print(f"Error migrating metrics: {e}")
    return count

def migrate_knowledge_graph(path: Path) -> int:
    """Stream large knowledge graph JSON to avoid OOM."""
    if not path.exists():
        return 0
    
    conn = get_conn()
    node_count, edge_count = 0, 0
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            nodes = data.get("nodes", {})
            edges = data.get("edges", [])
            
            with conn:
                for n_id, n_data in nodes.items():
                    conn.execute("INSERT OR IGNORE INTO graph_nodes (id, text, type) VALUES (?, ?, ?)", 
                                 (n_id, n_data.get("text", n_id), n_data.get("type", "CONCEPT")))
                    node_count += 1
                
                for e in edges:
                    conn.execute("INSERT OR IGNORE INTO graph_edges (from_node, to_node, edge_type, confidence) VALUES (?, ?, ?, ?)", 
                                 (e.get("from_node", e.get("from", "")), e.get("to_node", e.get("to", "")), e.get("edge_type", e.get("type", "")), float(e.get("confidence", 0.5))))
                    edge_count += 1
    except Exception as e:
        print(f"Error migrating graph: {e}")
    return node_count + edge_count

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--rollback", action="store_true", help="Delete the SQLite database.")
    args = parser.parse_args()

    if args.rollback:
        _rollback()
        return

    print("Initializing Database...")
    init_db()

    print("Migrating Lessons...")
    n = migrate_jsonl("lessons", LESSONS_DIR / "auto_lessons.jsonl", "lesson")
    print(f"  -> {n} lessons migrated.")

    print("Migrating Failures...")
    n = migrate_jsonl("failures", FAILS_DIR / "auto_fails.jsonl", "failure")
    print(f"  -> {n} failures migrated.")

    print("Migrating Patterns...")
    n = migrate_jsonl("patterns", PATTERNS_DIR / "auto_patterns.jsonl", "pattern")
    print(f"  -> {n} patterns migrated.")

    print("Migrating Facts...")
    n = migrate_json("facts", KNOWLEDGE_DIR / "auto_facts.json", "declarative")
    print(f"  -> {n} facts migrated.")

    print("Migrating Metrics...")
    n = migrate_retrieval_metrics(EVIDENCE_DIR / "Metrics" / "retrieval_metrics.json")
    print(f"  -> {n} metrics migrated.")

    print("Migrating Knowledge Graph...")
    n = migrate_knowledge_graph(GRAPH_FILE)
    print(f"  -> {n} graph elements migrated.")

    print("Migration Complete.")

if __name__ == "__main__":
    main()
