"""Firestore database integration."""
import firebase_admin
from firebase_admin import credentials, firestore
from typing import Optional, Dict, List, Any
from datetime import datetime
import os
import uuid
import json
from enum import Enum

# Initialize Firestore
if not firebase_admin._apps:
    try:
        cred_path = os.getenv("FIREBASE_ADMIN_CREDENTIALS_PATH")
        
        # Try to initialize from file path first
        if cred_path and os.path.exists(cred_path):
            cred = credentials.Certificate(cred_path)
            firebase_admin.initialize_app(cred)
        # Try environment variables (for Vercel/serverless)
        elif os.getenv("FIREBASE_PROJECT_ID") and os.getenv("FIREBASE_PRIVATE_KEY") and os.getenv("FIREBASE_CLIENT_EMAIL"):
            private_key = os.getenv("FIREBASE_PRIVATE_KEY")
            # Handle both escaped and unescaped newlines
            if private_key:
                private_key = private_key.replace("\\n", "\n")
            
            cred_dict = {
                "type": "service_account",
                "project_id": os.getenv("FIREBASE_PROJECT_ID"),
                "private_key_id": os.getenv("FIREBASE_PRIVATE_KEY_ID", ""),
                "private_key": private_key,
                "client_email": os.getenv("FIREBASE_CLIENT_EMAIL"),
                "client_id": os.getenv("FIREBASE_CLIENT_ID", ""),
                "auth_uri": "https://accounts.google.com/o/oauth2/auth",
                "token_uri": "https://oauth2.googleapis.com/token",
                "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
                "client_x509_cert_url": os.getenv("FIREBASE_CLIENT_X509_CERT_URL", "")
            }
            cred = credentials.Certificate(cred_dict)
            firebase_admin.initialize_app(cred)
    except Exception as e:
        print(f"❌ Error initializing Firebase: {e}")
        import traceback
        print(traceback.format_exc())
        # Don't raise - allow app to start even if Firebase fails (some endpoints might work)

# Get Firestore client (only if Firebase is initialized)
try:
    db = firestore.client() if firebase_admin._apps else None
except Exception as e:
    print(f"⚠️  Could not initialize Firestore client: {e}")
    db = None


# Collection names
COLLECTIONS = {
    'users': 'users',
    'organizations': 'organizations',
    'projects': 'projects',
    'org_members': 'org_members',
    'api_keys': 'api_keys',
    'audits': 'audits',
    'policy_versions': 'policy_versions'
}


def datetime_to_timestamp(dt: Optional[datetime]) -> Optional[Any]:
    """Convert datetime to Firestore timestamp."""
    if dt is None:
        return None
    # Firestore accepts datetime objects directly
    return dt


def timestamp_to_datetime(ts: Any) -> Optional[datetime]:
    """Convert Firestore timestamp to datetime."""
    if ts is None:
        return None
    # Firestore returns timestamps that can be used as datetime
    if hasattr(ts, 'timestamp'):
        return datetime.fromtimestamp(ts.timestamp())
    return ts


def to_dict(data: Any, exclude_none: bool = False) -> Dict:
    """Convert data to dictionary, handling enums and datetimes."""
    if isinstance(data, dict):
        result = {}
        for k, v in data.items():
            if exclude_none and v is None:
                continue
            if isinstance(v, Enum):
                result[k] = v.value
            elif isinstance(v, datetime):
                result[k] = v
            elif isinstance(v, dict):
                result[k] = to_dict(v, exclude_none)
            else:
                result[k] = v
        return result
    elif hasattr(data, '__dict__'):
        result = {}
        for k, v in data.__dict__.items():
            if k.startswith('_'):
                continue
            if exclude_none and v is None:
                continue
            if isinstance(v, Enum):
                result[k] = v.value
            elif isinstance(v, datetime):
                result[k] = v
            elif isinstance(v, dict):
                result[k] = to_dict(v, exclude_none)
            else:
                result[k] = v
        return result
    return data


# User operations
def get_user(user_id: str) -> Optional[Dict]:
    """Get user by ID."""
    if not db:
        return None
    doc = db.collection(COLLECTIONS['users']).document(user_id).get()
    return doc.to_dict() if doc.exists else None


def get_user_by_email(email: str) -> Optional[Dict]:
    """Get user by email."""
    if not db:
        return None
    users = db.collection(COLLECTIONS['users']).where('email', '==', email).limit(1).stream()
    for doc in users:
        return doc.to_dict()
    return None


def create_user(user_data: Dict) -> Dict:
    """Create a new user."""
    if not db:
        raise Exception("Firestore not initialized")
    
    user_id = user_data.get('id') or str(uuid.uuid4())
    user_data['id'] = user_id
    user_data['created_at'] = datetime.utcnow()
    
    # Convert enums to strings
    if 'subscription_plan' in user_data and isinstance(user_data['subscription_plan'], Enum):
        user_data['subscription_plan'] = user_data['subscription_plan'].value
    
    db.collection(COLLECTIONS['users']).document(user_id).set(to_dict(user_data))
    return user_data


def update_user(user_id: str, updates: Dict) -> Dict:
    """Update user."""
    if not db:
        raise Exception("Firestore not initialized")
    
    # Convert enums to strings
    if 'subscription_plan' in updates and isinstance(updates['subscription_plan'], Enum):
        updates['subscription_plan'] = updates['subscription_plan'].value
    
    db.collection(COLLECTIONS['users']).document(user_id).update(to_dict(updates, exclude_none=True))
    return get_user(user_id)


# Organization operations
def get_organization(org_id: str) -> Optional[Dict]:
    """Get organization by ID."""
    if not db:
        return None
    doc = db.collection(COLLECTIONS['organizations']).document(org_id).get()
    return doc.to_dict() if doc.exists else None


def get_organizations_by_user(user_id: str) -> List[Dict]:
    """Get all organizations for a user."""
    if not db:
        return []
    
    # Get memberships
    memberships = db.collection(COLLECTIONS['org_members']).where('user_id', '==', user_id).stream()
    org_ids = [m.to_dict()['organization_id'] for m in memberships]
    
    if not org_ids:
        return []
    
    # Get organizations
    organizations = []
    for org_id in org_ids:
        org = get_organization(org_id)
        if org:
            organizations.append(org)
    
    return organizations


def create_organization(org_data: Dict) -> Dict:
    """Create a new organization."""
    if not db:
        raise Exception("Firestore not initialized")
    
    org_id = org_data.get('id') or str(uuid.uuid4())
    org_data['id'] = org_id
    org_data['created_at'] = datetime.utcnow()
    
    # Convert enums to strings
    if 'plan' in org_data and isinstance(org_data['plan'], Enum):
        org_data['plan'] = org_data['plan'].value
    
    db.collection(COLLECTIONS['organizations']).document(org_id).set(to_dict(org_data))
    return org_data


def update_organization(org_id: str, updates: Dict) -> Dict:
    """Update organization."""
    if not db:
        raise Exception("Firestore not initialized")
    
    # Convert enums to strings
    if 'plan' in updates and isinstance(updates['plan'], Enum):
        updates['plan'] = updates['plan'].value
    
    db.collection(COLLECTIONS['organizations']).document(org_id).update(to_dict(updates, exclude_none=True))
    return get_organization(org_id)


def get_organization_by_stripe_customer(stripe_customer_id: str) -> Optional[Dict]:
    """Get organization by Stripe customer ID."""
    if not db:
        return None
    orgs = db.collection(COLLECTIONS['organizations']).where('stripe_customer_id', '==', stripe_customer_id).limit(1).stream()
    for doc in orgs:
        data = doc.to_dict()
        data['id'] = doc.id  # Ensure ID is included
        return data
    return None


# Project operations
def get_project(project_id: str) -> Optional[Dict]:
    """Get project by ID."""
    if not db:
        return None
    doc = db.collection(COLLECTIONS['projects']).document(project_id).get()
    return doc.to_dict() if doc.exists else None


def get_projects_by_organization(org_id: str) -> List[Dict]:
    """Get all projects for an organization."""
    if not db:
        return []
    projects = db.collection(COLLECTIONS['projects']).where('organization_id', '==', org_id).stream()
    return [doc.to_dict() for doc in projects]


def create_project(project_data: Dict) -> Dict:
    """Create a new project."""
    if not db:
        raise Exception("Firestore not initialized")
    
    project_id = project_data.get('id') or str(uuid.uuid4())
    project_data['id'] = project_id
    project_data['created_at'] = datetime.utcnow()
    
    db.collection(COLLECTIONS['projects']).document(project_id).set(to_dict(project_data))
    return project_data


def update_project(project_id: str, updates: Dict) -> Dict:
    """Update project."""
    if not db:
        raise Exception("Firestore not initialized")
    
    db.collection(COLLECTIONS['projects']).document(project_id).update(to_dict(updates, exclude_none=True))
    return get_project(project_id)


# OrgMember operations
def create_org_member(member_data: Dict) -> Dict:
    """Create organization membership."""
    if not db:
        raise Exception("Firestore not initialized")
    
    member_id = member_data.get('id') or str(uuid.uuid4())
    member_data['id'] = member_id
    member_data['created_at'] = datetime.utcnow()
    
    # Convert enums to strings
    if 'role' in member_data and isinstance(member_data['role'], Enum):
        member_data['role'] = member_data['role'].value
    
    db.collection(COLLECTIONS['org_members']).document(member_id).set(to_dict(member_data))
    return member_data


def get_org_members(org_id: str) -> List[Dict]:
    """Get all members of an organization."""
    if not db:
        return []
    members = db.collection(COLLECTIONS['org_members']).where('organization_id', '==', org_id).stream()
    results = []
    for doc in members:
        data = doc.to_dict()
        data['id'] = doc.id  # Ensure ID is included
        results.append(data)
    return results


def get_org_member(user_id: str, org_id: str) -> Optional[Dict]:
    """Get specific org member."""
    if not db:
        return None
    members = db.collection(COLLECTIONS['org_members']).where('user_id', '==', user_id).where('organization_id', '==', org_id).limit(1).stream()
    for doc in members:
        return doc.to_dict()
    return None


# APIKey operations
def get_api_key(key_hash: str) -> Optional[Dict]:
    """Get API key by hash."""
    if not db:
        return None
    keys = db.collection(COLLECTIONS['api_keys']).where('key_hash', '==', key_hash).limit(1).stream()
    for doc in keys:
        return doc.to_dict()
    return None


def get_api_keys_by_organization(org_id: str, project_id: Optional[str] = None) -> List[Dict]:
    """Get API keys for organization/project."""
    if not db:
        return []
    
    query = db.collection(COLLECTIONS['api_keys']).where('organization_id', '==', org_id)
    if project_id:
        query = query.where('project_id', '==', project_id)
    # If project_id is None, don't filter by project_id - return all keys for the org
    
    keys = query.stream()
    return [doc.to_dict() for doc in keys]


def create_api_key(key_data: Dict) -> Dict:
    """Create API key."""
    if not db:
        raise Exception("Firestore not initialized")
    
    key_id = key_data.get('id') or str(uuid.uuid4())
    key_data['id'] = key_id
    key_data['created_at'] = datetime.utcnow()
    if 'rate_limit_reset' not in key_data:
        key_data['rate_limit_reset'] = datetime.utcnow()
    
    db.collection(COLLECTIONS['api_keys']).document(key_id).set(to_dict(key_data))
    return key_data


def get_api_key_by_id(key_id: str) -> Optional[Dict]:
    """Get API key by document ID."""
    if not db:
        return None
    doc = db.collection(COLLECTIONS['api_keys']).document(key_id).get()
    if doc.exists:
        data = doc.to_dict()
        data['id'] = doc.id  # Ensure ID is included
        return data
    return None


def update_api_key(key_id: str, updates: Dict) -> Dict:
    """Update API key."""
    if not db:
        raise Exception("Firestore not initialized")
    
    db.collection(COLLECTIONS['api_keys']).document(key_id).update(to_dict(updates, exclude_none=True))
    doc = db.collection(COLLECTIONS['api_keys']).document(key_id).get()
    return doc.to_dict() if doc.exists else None


# Audit operations
def get_audit(audit_id: str) -> Optional[Dict]:
    """Get audit by ID."""
    if not db:
        return None
    doc = db.collection(COLLECTIONS['audits']).document(audit_id).get()
    return doc.to_dict() if doc.exists else None


def get_audits_by_project(project_id: str, limit: int = 50, offset: int = 0) -> List[Dict]:
    """Get audits for a project."""
    if not db:
        return []
    
    # Firestore doesn't support offset directly, we'll use start_after for pagination
    # For now, just get the first page
    query = db.collection(COLLECTIONS['audits']).where('project_id', '==', project_id).order_by('created_at', direction=firestore.Query.DESCENDING).limit(limit)
    audits = query.stream()
    results = [doc.to_dict() for doc in audits]
    # Apply offset manually (not ideal but works for small datasets)
    return results[offset:]


def create_audit(audit_data: Dict) -> Dict:
    """Create audit (append-only)."""
    if not db:
        raise Exception("Firestore not initialized")
    
    audit_id = audit_data.get('id')
    if not audit_id:
        raise Exception("audit_id is required")
    
    # Check if audit already exists (append-only)
    existing = get_audit(audit_id)
    if existing:
        raise Exception(f"Audit {audit_id} already exists (append-only storage)")
    
    if 'created_at' not in audit_data:
        audit_data['created_at'] = datetime.utcnow()
    
    # Store report_json as string if it's a dict
    if 'report_json' in audit_data and isinstance(audit_data['report_json'], dict):
        audit_data['report_json'] = json.dumps(audit_data['report_json'])
    
    db.collection(COLLECTIONS['audits']).document(audit_id).set(to_dict(audit_data))
    return audit_data


# Policy Version operations
def get_policy_version(project_id: str, version: int) -> Optional[Dict]:
    """Get policy version by project ID and version number."""
    if not db:
        return None
    versions = db.collection(COLLECTIONS['policy_versions']).where('project_id', '==', project_id).where('version', '==', version).limit(1).stream()
    for doc in versions:
        data = doc.to_dict()
        data['id'] = doc.id
        return data
    return None


def get_policy_versions(project_id: str) -> List[Dict]:
    """Get all policy versions for a project, ordered by version descending."""
    if not db:
        return []
    versions = db.collection(COLLECTIONS['policy_versions']).where('project_id', '==', project_id).order_by('version', direction=firestore.Query.DESCENDING).stream()
    results = []
    for doc in versions:
        data = doc.to_dict()
        data['id'] = doc.id
        results.append(data)
    return results


def create_policy_version(version_data: Dict) -> Dict:
    """Create a new policy version."""
    if not db:
        raise Exception("Firestore not initialized")
    
    version_id = version_data.get('id') or str(uuid.uuid4())
    version_data['id'] = version_id
    if 'created_at' not in version_data:
        version_data['created_at'] = datetime.utcnow()
    
    db.collection(COLLECTIONS['policy_versions']).document(version_id).set(to_dict(version_data))
    return version_data


def get_user_by_stripe_customer(stripe_customer_id: str) -> Optional[Dict]:
    """Get user by Stripe customer ID."""
    if not db:
        return None
    users = db.collection(COLLECTIONS['users']).where('stripe_customer_id', '==', stripe_customer_id).limit(1).stream()
    for doc in users:
        user_data = doc.to_dict()
        user_data['id'] = doc.id  # Ensure ID is included
        return user_data
    return None

