import json
from .logger_util import logger

def save_schema_to_json(report, output_path="schema_map.json"):
    """
    Transforms the raw inference report into a clean, human-readable JSON
    mapping Primary Keys and Foreign Keys with confidence scores.
    """
    formatted_output = {
        "metadata": {
            "generated_at": "timestamp_here", # You can add datetime
            "stats": {
                "tables_analyzed": 0,
                "relationships_found": 0
            }
        },
        "tables": {},
        "relationships": []
    }

    # 1. Process Tables (Primary Keys)
    # Exclude the 'relations' key which sits at the root
    table_names = [k for k in report.keys() if k != 'relations']
    formatted_output["metadata"]["stats"]["tables_analyzed"] = len(table_names)

    for table in table_names:
        pks_data = report[table].get('pks', {})
        pk_list = []
        
        for col_name, attrs in pks_data.items():
            pk_list.append({
                "column": col_name,
                "type": "surrogate" if attrs.get("is_surrogate") else "natural"
            })

        formatted_output["tables"][table] = {
            "primary_keys": pk_list
        }

    # 2. Process Relationships (Aggregate Duplicates)
    # Multiple strategies might find the same link. We want to list them together.
    # Key = "source_table.col -> target_table.col"
    rel_map = {}

    for rel in report.get('relations', []):
        # Create a unique key for the relationship
        rel_key = f"{rel['source_table']}.{rel['source_col']}->{rel['target_table']}.{rel['target_col']}"

        if rel_key not in rel_map:
            rel_map[rel_key] = {
                "source": {
                    "table": rel['source_table'],
                    "column": rel['source_col']
                },
                "target": {
                    "table": rel['target_table'],
                    "column": rel['target_col']
                },
                "evidence": [],     # List of strategies that found this
                "overall_confidence": "Low" 
            }

        # Add the strategy evidence
        rel_map[rel_key]["evidence"].append({
            "strategy": rel['strategy'],
            "confidence": rel.get('confidence', 'Unknown')
        })

    # 3. Calculate Final Confidence
    # If explicit constraint -> 100%
    # If Naming + Containment -> High
    # If Naming only -> Medium/High
    for rel_key, data in rel_map.items():
        strategies = [e['strategy'] for e in data['evidence']]
        
        if 'ExistingConstraint' in strategies:
            data['overall_confidence'] = "Certain (Database Constraint)"
        elif 'ValueContainment' in strategies:
            data['overall_confidence'] = "Very High (Validated Data)"
        elif 'NamingConvention' in strategies:
            data['overall_confidence'] = "High (Naming Match)"
        
        formatted_output["relationships"].append(data)

    formatted_output["metadata"]["stats"]["relationships_found"] = len(formatted_output["relationships"])

    # 4. Save to File
    try:
        with open(output_path, 'w') as f:
            json.dump(formatted_output, f, indent=4)
        logger.info(f"✅ Schema map saved successfully to: {output_path}")
        return formatted_output
    except Exception as e:
        logger.error(f"Failed to save JSON report: {e}")
        return {}