#!/usr/bin/env python3
"""
Migration script for importing YAML data into Reality Check.

Handles:
- Domain migration (VALUE/DIST → ECON, SOCIAL → SOC)
- Claim ID renumbering to avoid collisions
- Confidence → Credence field rename
- Embedding generation for all records
- Prediction extraction from predictions.md
"""

from __future__ import annotations

import argparse
import re
import sys
from collections import defaultdict
from datetime import date
from pathlib import Path
from typing import Any, Optional

import yaml

# Add scripts directory to path for imports
sys.path.insert(0, str(Path(__file__).parent))

from db import (
    DOMAIN_MIGRATION,
    VALID_DOMAINS,
    add_chain,
    add_claim,
    add_prediction,
    add_source,
    drop_tables,
    get_db,
    init_tables,
    get_stats,
)


def load_yaml(path: Path) -> dict:
    """Load a YAML file."""
    with open(path) as f:
        return yaml.safe_load(f)


def migrate_domain(domain: str) -> str:
    """Migrate old domain to new domain."""
    return DOMAIN_MIGRATION.get(domain, domain)


def build_id_mapping(claims: dict[str, dict]) -> dict[str, str]:
    """
    Build mapping from old claim IDs to new claim IDs.

    Migration rules:
    - VALUE-* → ECON-* (renumbered)
    - DIST-* → ECON-* (continue from VALUE max)
    - SOCIAL-* → SOC-*
    - Others stay the same
    """
    mapping = {}

    # Track new ECON IDs by year
    econ_counters: dict[int, int] = defaultdict(int)

    # First pass: count existing ECON claims to set starting point
    # (In case there are any direct ECON claims, which there shouldn't be in v0)

    # Sort claims by ID for deterministic ordering
    sorted_claims = sorted(claims.items())

    for old_id, claim in sorted_claims:
        parts = old_id.split("-")
        if len(parts) != 3:
            mapping[old_id] = old_id  # Invalid ID, keep as-is
            continue

        old_domain, year_str, num_str = parts
        year = int(year_str)

        new_domain = migrate_domain(old_domain)

        if old_domain in ("VALUE", "DIST"):
            # These merge into ECON, need new numbers
            econ_counters[year] += 1
            new_num = econ_counters[year]
            new_id = f"ECON-{year:04d}-{new_num:03d}"
            mapping[old_id] = new_id
        elif old_domain == "SOCIAL":
            # Simple rename, keep number
            new_id = f"SOC-{year_str}-{num_str}"
            mapping[old_id] = new_id
        else:
            # Keep original ID
            mapping[old_id] = old_id

    return mapping


def update_claim_references(claim: dict, id_mapping: dict[str, str]) -> dict:
    """Update all claim ID references in a claim using the mapping."""
    claim = claim.copy()

    # Update relationship fields
    for field in ["supports", "contradicts", "depends_on", "modified_by"]:
        if field in claim and claim[field]:
            claim[field] = [id_mapping.get(ref, ref) for ref in claim[field]]

    # Update part_of_chain (chain IDs don't change)
    # No change needed

    return claim


def migrate_claim(old_id: str, claim: dict, id_mapping: dict[str, str]) -> dict:
    """Convert a claim from YAML format to DB format."""
    new_id = id_mapping.get(old_id, old_id)
    new_domain = migrate_domain(claim.get("domain", ""))

    # Update references
    claim = update_claim_references(claim, id_mapping)

    return {
        "id": new_id,
        "text": claim.get("text", ""),
        "type": claim.get("type", "[T]"),
        "domain": new_domain,
        "evidence_level": claim.get("evidence_level", "E4"),
        "credence": float(claim.get("confidence", 0.5)),  # Rename: confidence → credence

        # v1.0 additions (empty for migrated claims)
        "operationalization": None,
        "assumptions": [],
        "falsifiers": [],

        # Provenance
        "source_ids": claim.get("source_ids", []),
        "first_extracted": claim.get("first_extracted", str(date.today())),
        "extracted_by": claim.get("extracted_by", "migration"),

        # Relationships (already updated)
        "supports": claim.get("supports", []),
        "contradicts": claim.get("contradicts", []),
        "depends_on": claim.get("depends_on", []),
        "modified_by": claim.get("modified_by", []),
        "part_of_chain": claim.get("part_of_chain", ""),

        # Versioning
        "version": claim.get("version", 1),
        "last_updated": claim.get("last_updated", str(date.today())),
        "notes": claim.get("notes", None),

        # Embedding will be generated by add_claim
        "embedding": None,
    }


def migrate_source(source_id: str, source: dict, id_mapping: dict[str, str]) -> dict:
    """Convert a source from YAML format to DB format."""
    # Update claims_extracted with new IDs
    claims_extracted = source.get("claims_extracted", []) or []
    claims_extracted = [id_mapping.get(cid, cid) for cid in claims_extracted]

    # Update domains
    domains = source.get("domains", []) or []
    domains = list(set(migrate_domain(d) for d in domains))

    return {
        "id": source_id,
        "type": source.get("type", "ARTICLE"),
        "title": source.get("title", ""),
        "author": source.get("author", []),
        "year": source.get("year", 0),
        "url": source.get("url", ""),
        "doi": source.get("doi", None),
        "accessed": source.get("accessed", None),
        "reliability": float(source.get("reliability", 0.5)) if source.get("reliability") else None,
        "bias_notes": source.get("bias_notes", None),
        "claims_extracted": claims_extracted,
        "analysis_file": source.get("analysis_file", None),
        "topics": source.get("topics", []) or [],
        "domains": domains,
        "status": source.get("status", "cataloged"),

        # Embedding will be generated by add_source
        "embedding": None,
    }


def migrate_chain(chain_id: str, chain: dict, id_mapping: dict[str, str]) -> dict:
    """Convert a chain from YAML format to DB format."""
    # Update claim IDs in the chain
    claims = chain.get("claims", []) or []
    claims = [id_mapping.get(cid, cid) for cid in claims]

    # Update weakest_link reference if it contains a claim ID
    weakest_link = chain.get("weakest_link", "")
    if weakest_link:
        for old_id, new_id in id_mapping.items():
            if old_id in weakest_link:
                weakest_link = weakest_link.replace(old_id, new_id)

    return {
        "id": chain_id,
        "name": chain.get("name", ""),
        "thesis": chain.get("thesis", ""),
        "credence": float(chain.get("confidence", 0.5)),  # Rename: confidence → credence
        "claims": claims,
        "analysis_file": chain.get("analysis_file", None),
        "weakest_link": weakest_link,
        "scoring_method": "MIN",  # Default scoring method

        # Embedding will be generated by add_chain
        "embedding": None,
    }


def parse_predictions_md(predictions_path: Path, id_mapping: dict[str, str]) -> list[dict]:
    """Parse predictions.md and extract prediction records."""
    predictions = []

    if not predictions_path.exists():
        return predictions

    text = predictions_path.read_text()

    # Parse prediction blocks
    # Format: ### [Title]\n- **Claim ID**: ...\n- **Source**: ...
    blocks = re.split(r'\n### ', text)

    for block in blocks[1:]:  # Skip header
        prediction = {}

        # Extract fields using regex
        claim_match = re.search(r'\*\*Claim ID\*\*:\s*([A-Z]+-\d{4}-\d{3})', block)
        source_match = re.search(r'\*\*Source\*\*:\s*([a-z0-9][a-z0-9\-]*)', block)
        status_match = re.search(r'\*\*Status\*\*:\s*(\[P[+~→?←!∅-]\])', block)
        date_made_match = re.search(r'\*\*Date Made\*\*:\s*(\d{4}-\d{2}-\d{2})', block)
        target_match = re.search(r'\*\*Target Date\*\*:\s*([^\n]+)', block)
        falsification_match = re.search(r'\*\*Falsification Criteria\*\*:\s*([^\n]+)', block)
        verification_match = re.search(r'\*\*Verification Criteria\*\*:\s*([^\n]+)', block)
        last_eval_match = re.search(r'\*\*Last Evaluated\*\*:\s*(\d{4}-\d{2}-\d{2})', block)

        if claim_match:
            old_claim_id = claim_match.group(1)
            new_claim_id = id_mapping.get(old_claim_id, old_claim_id)

            prediction = {
                "claim_id": new_claim_id,
                "source_id": source_match.group(1) if source_match else "",
                "date_made": date_made_match.group(1) if date_made_match else None,
                "target_date": target_match.group(1).strip() if target_match else None,
                "falsification_criteria": falsification_match.group(1).strip() if falsification_match else None,
                "verification_criteria": verification_match.group(1).strip() if verification_match else None,
                "status": status_match.group(1) if status_match else "[P?]",
                "last_evaluated": last_eval_match.group(1) if last_eval_match else None,
                "evidence_updates": None,  # Would need more parsing
            }
            predictions.append(prediction)

    return predictions


def run_migration(
    source_repo: Path,
    db_path: Optional[Path] = None,
    dry_run: bool = False,
    verbose: bool = False,
) -> dict:
    """
    Run the full migration from YAML to LanceDB.

    Args:
        source_repo: Path to the source repository with YAML data
        db_path: Optional override for database path
        dry_run: If True, don't actually write to DB
        verbose: If True, print progress

    Returns:
        Migration statistics
    """
    stats = {
        "claims_migrated": 0,
        "sources_migrated": 0,
        "chains_migrated": 0,
        "predictions_migrated": 0,
        "id_mappings": {},
        "errors": [],
    }

    # Load source YAML files
    claims_path = source_repo / "claims" / "registry.yaml"
    sources_path = source_repo / "reference" / "sources.yaml"
    predictions_path = source_repo / "tracking" / "predictions.md"

    if not claims_path.exists():
        stats["errors"].append(f"Claims registry not found: {claims_path}")
        return stats

    if not sources_path.exists():
        stats["errors"].append(f"Sources registry not found: {sources_path}")
        return stats

    if verbose:
        print(f"Loading claims from {claims_path}")
    claims_data = load_yaml(claims_path)

    if verbose:
        print(f"Loading sources from {sources_path}")
    sources_data = load_yaml(sources_path)

    # Extract data
    claims = claims_data.get("claims", {})
    chains = claims_data.get("chains", {})
    sources = sources_data.get("sources", {})

    if verbose:
        print(f"Found {len(claims)} claims, {len(chains)} chains, {len(sources)} sources")

    # Build ID mapping
    id_mapping = build_id_mapping(claims)
    stats["id_mappings"] = id_mapping

    if verbose:
        changed = {k: v for k, v in id_mapping.items() if k != v}
        print(f"ID mappings ({len(changed)} changed):")
        for old, new in sorted(changed.items()):
            print(f"  {old} → {new}")

    if dry_run:
        if verbose:
            print("\n[DRY RUN] Would migrate:")
            print(f"  {len(claims)} claims")
            print(f"  {len(sources)} sources")
            print(f"  {len(chains)} chains")
        return stats

    # Initialize database
    if verbose:
        print(f"\nInitializing database...")

    db = get_db(db_path)
    drop_tables(db)  # Start fresh
    init_tables(db)

    # Migrate sources first (claims reference them)
    if verbose:
        print(f"\nMigrating {len(sources)} sources...")

    for source_id, source in sources.items():
        try:
            migrated = migrate_source(source_id, source, id_mapping)
            add_source(migrated, db)
            stats["sources_migrated"] += 1
            if verbose:
                print(f"  ✓ {source_id}")
        except Exception as e:
            stats["errors"].append(f"Source {source_id}: {e}")
            if verbose:
                print(f"  ✗ {source_id}: {e}")

    # Migrate claims
    if verbose:
        print(f"\nMigrating {len(claims)} claims...")

    for old_id, claim in claims.items():
        try:
            migrated = migrate_claim(old_id, claim, id_mapping)
            add_claim(migrated, db)
            stats["claims_migrated"] += 1
            new_id = id_mapping.get(old_id, old_id)
            if verbose:
                if old_id != new_id:
                    print(f"  ✓ {old_id} → {new_id}")
                else:
                    print(f"  ✓ {old_id}")
        except Exception as e:
            stats["errors"].append(f"Claim {old_id}: {e}")
            if verbose:
                print(f"  ✗ {old_id}: {e}")

    # Migrate chains
    if verbose:
        print(f"\nMigrating {len(chains)} chains...")

    for chain_id, chain in chains.items():
        try:
            migrated = migrate_chain(chain_id, chain, id_mapping)
            add_chain(migrated, db)
            stats["chains_migrated"] += 1
            if verbose:
                print(f"  ✓ {chain_id}")
        except Exception as e:
            stats["errors"].append(f"Chain {chain_id}: {e}")
            if verbose:
                print(f"  ✗ {chain_id}: {e}")

    # Migrate predictions
    if verbose:
        print(f"\nParsing predictions from {predictions_path}...")

    predictions = parse_predictions_md(predictions_path, id_mapping)

    if verbose:
        print(f"Found {len(predictions)} predictions")

    for pred in predictions:
        try:
            add_prediction(pred, db)
            stats["predictions_migrated"] += 1
            if verbose:
                print(f"  ✓ {pred['claim_id']}")
        except Exception as e:
            stats["errors"].append(f"Prediction {pred.get('claim_id', '?')}: {e}")
            if verbose:
                print(f"  ✗ {pred.get('claim_id', '?')}: {e}")

    # Print summary
    if verbose:
        print(f"\n{'='*50}")
        print("Migration Summary:")
        print(f"  Claims: {stats['claims_migrated']}")
        print(f"  Sources: {stats['sources_migrated']}")
        print(f"  Chains: {stats['chains_migrated']}")
        print(f"  Predictions: {stats['predictions_migrated']}")
        if stats["errors"]:
            print(f"  Errors: {len(stats['errors'])}")
            for err in stats["errors"]:
                print(f"    - {err}")

        # Verify with stats
        db_stats = get_stats(db)
        print(f"\nDatabase stats after migration:")
        for table, count in db_stats.items():
            print(f"  {table}: {count}")

    return stats


def main():
    parser = argparse.ArgumentParser(
        description="Migrate YAML data to LanceDB"
    )
    parser.add_argument(
        "source_repo",
        type=Path,
        help="Path to source repository with YAML data",
    )
    parser.add_argument(
        "--db-path",
        type=Path,
        default=None,
        help="Override database path (default: data/realitycheck.lance)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Show what would be migrated without writing",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Verbose output",
    )

    args = parser.parse_args()

    if not args.source_repo.exists():
        print(f"Error: Source repo not found: {args.source_repo}", file=sys.stderr)
        sys.exit(1)

    stats = run_migration(
        source_repo=args.source_repo,
        db_path=args.db_path,
        dry_run=args.dry_run,
        verbose=args.verbose,
    )

    if stats["errors"]:
        sys.exit(1)


if __name__ == "__main__":
    main()
