use crate::models::{
    Tenant, DataSource, ExternalEntity, CanonicalEntity, IdentityLink, MatchRule, ManualOverride,
    User, BatchRun, AuditEvent, RefreshToken, SsoConfig, AuditStream, RetentionPolicy, MatchResult, MatchAudit,
    SurvivorshipPolicy, SchemaMapping, ApiKey, IdentityPlanModel
};
use crate::error::Result;
use sqlx::PgPool;
use redis::AsyncCommands;
use uuid::Uuid;
use async_trait::async_trait;
use sha2::Digest;

#[async_trait]
pub trait PgConnectionExt {
    async fn set_tenant_context(&mut self, tenant_id: Uuid) -> sqlx::Result<()>;
}

#[async_trait]
impl PgConnectionExt for sqlx::PgConnection {
    async fn set_tenant_context(&mut self, tenant_id: Uuid) -> sqlx::Result<()> {
        sqlx::query("SELECT set_config('request.tenant_id', $1, true)")
            .bind(tenant_id.to_string())
            .execute(self)
            .await?;
        Ok(())
    }
}

#[async_trait]
impl<'c> PgConnectionExt for sqlx::Transaction<'c, sqlx::Postgres> {
    async fn set_tenant_context(&mut self, tenant_id: Uuid) -> sqlx::Result<()> {
        sqlx::query("SELECT set_config('request.tenant_id', $1, true)")
            .bind(tenant_id.to_string())
            .execute(&mut **self)
            .await?;
        Ok(())
    }
}

#[async_trait]
pub trait TenantRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid) -> Result<Option<Tenant>>;
    async fn find_by_slug(&self, slug: &str) -> Result<Option<Tenant>>;
    async fn create(&self, name: &str, slug: &str, parent_id: Option<Uuid>, tier: &str, usage_limits: serde_json::Value) -> Result<Tenant>;
    async fn list_children(&self, parent_id: Uuid) -> Result<Vec<Tenant>>;
    async fn list_all(&self) -> Result<Vec<Tenant>>;
    async fn update_tier(&self, id: Uuid, tier: &str, usage_limits: serde_json::Value) -> Result<Tenant>;
    async fn update_settings(&self, id: Uuid, settings: serde_json::Value) -> Result<Tenant>;
    async fn update(&self, tenant: Tenant) -> Result<Tenant>;
}

pub struct SqlxTenantRepository { pool: PgPool }
impl SqlxTenantRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl TenantRepository for SqlxTenantRepository {
    async fn find_by_id(&self, id: Uuid) -> Result<Option<Tenant>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(id).await?; // A tenant can always see itself
        let res = sqlx::query_as::<_, Tenant>("SELECT * FROM tenants WHERE id = $1").bind(id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn find_by_slug(&self, slug: &str) -> Result<Option<Tenant>> {
        // This is tricky because we don't know the ID yet.
        // For slug lookup, we might need a bypass or use a system context.
        // For now, assume it's used in contexts where RLS is not yet active or we use the pool directly.
        Ok(sqlx::query_as::<_, Tenant>("SELECT * FROM tenants WHERE slug = $1").bind(slug).fetch_optional(&self.pool).await?)
    }
    async fn create(&self, name: &str, slug: &str, parent_id: Option<Uuid>, tier: &str, usage_limits: serde_json::Value) -> Result<Tenant> {
        Ok(sqlx::query_as::<_, Tenant>("INSERT INTO tenants (name, slug, parent_id, tier, usage_limits) VALUES ($1, $2, $3, $4, $5) RETURNING *")
            .bind(name)
            .bind(slug)
            .bind(parent_id)
            .bind(tier)
            .bind(usage_limits)
            .fetch_one(&self.pool)
            .await?)
    }
    async fn list_children(&self, parent_id: Uuid) -> Result<Vec<Tenant>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(parent_id).await?;
        let res = sqlx::query_as::<_, Tenant>("SELECT * FROM tenants WHERE parent_id = $1").bind(parent_id).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_all(&self) -> Result<Vec<Tenant>> {
        Ok(sqlx::query_as::<_, Tenant>("SELECT * FROM tenants ORDER BY created_at DESC").fetch_all(&self.pool).await?)
    }
    async fn update_tier(&self, id: Uuid, tier: &str, usage_limits: serde_json::Value) -> Result<Tenant> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(id).await?;
        let res = sqlx::query_as::<_, Tenant>("UPDATE tenants SET tier = $2, usage_limits = $3 WHERE id = $1 RETURNING *").bind(id).bind(tier).bind(usage_limits).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update_settings(&self, id: Uuid, settings: serde_json::Value) -> Result<Tenant> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(id).await?;
        let res = sqlx::query_as::<_, Tenant>("UPDATE tenants SET settings = $2 WHERE id = $1 RETURNING *").bind(id).bind(settings).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update(&self, tenant: Tenant) -> Result<Tenant> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant.id).await?;
        let res = sqlx::query_as::<_, Tenant>(
            "UPDATE tenants SET name = $2, slug = $3, settings = $4, tier = $5, usage_limits = $6, billing_id = $7 WHERE id = $1 RETURNING *"
        )
        .bind(tenant.id)
        .bind(tenant.name)
        .bind(tenant.slug)
        .bind(tenant.settings)
        .bind(tenant.tier)
        .bind(tenant.usage_limits)
        .bind(tenant.billing_id)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
}

#[async_trait]
pub trait DataSourceRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<DataSource>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<DataSource>>;
    async fn create(&self, tenant_id: Uuid, name: &str, source_type: &str, config: serde_json::Value) -> Result<DataSource>;
    async fn update(&self, id: Uuid, tenant_id: Uuid, name: &str, config: serde_json::Value) -> Result<DataSource>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxDataSourceRepository { pool: PgPool }
impl SqlxDataSourceRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl DataSourceRepository for SqlxDataSourceRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<DataSource>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, DataSource>("SELECT * FROM data_sources WHERE id = $1").bind(id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<DataSource>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, DataSource>("SELECT * FROM data_sources").fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create(&self, tenant_id: Uuid, name: &str, source_type: &str, config: serde_json::Value) -> Result<DataSource> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, DataSource>("INSERT INTO data_sources (tenant_id, name, source_type, config) VALUES ($1, $2, $3, $4) RETURNING *").bind(tenant_id).bind(name).bind(source_type).bind(config).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update(&self, id: Uuid, tenant_id: Uuid, name: &str, config: serde_json::Value) -> Result<DataSource> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, DataSource>("UPDATE data_sources SET name = $2, config = $3 WHERE id = $1 AND tenant_id = $4 RETURNING *")
            .bind(id).bind(name).bind(config).bind(tenant_id)
            .fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        res.ok_or_else(|| crate::error::CannonError::NotFound(format!("Source {} not found or access denied", id)).into())
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let result = sqlx::query("DELETE FROM data_sources WHERE id = $1 AND tenant_id = $2").bind(id).bind(tenant_id).execute(&mut *tx).await?;
        tx.commit().await?;
        if result.rows_affected() == 0 {
            return Err(crate::error::CannonError::NotFound(format!("Source {} not found or access denied", id)).into());
        }
        Ok(())
    }
}

#[async_trait]
pub trait ExternalEntityRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<ExternalEntity>>;
    async fn list_by_tenant(&self, tenant_id: Uuid, entity_type: &str, limit: i64, offset: i64) -> Result<Vec<ExternalEntity>>;
    async fn list_by_source(&self, tenant_id: Uuid, source_id: Uuid, limit: i64) -> Result<Vec<ExternalEntity>>;
    async fn upsert(&self, entity: ExternalEntity) -> Result<ExternalEntity>;
}

pub struct SqlxExternalEntityRepository { pool: PgPool }
impl SqlxExternalEntityRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl ExternalEntityRepository for SqlxExternalEntityRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<ExternalEntity>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ExternalEntity>("SELECT * FROM external_entities WHERE id = $1")
            .bind(id)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid, entity_type: &str, limit: i64, offset: i64) -> Result<Vec<ExternalEntity>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ExternalEntity>("SELECT * FROM external_entities WHERE entity_type = $1 LIMIT $2 OFFSET $3")
            .bind(entity_type).bind(limit).bind(offset).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_source(&self, tenant_id: Uuid, source_id: Uuid, limit: i64) -> Result<Vec<ExternalEntity>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ExternalEntity>("SELECT * FROM external_entities WHERE data_source_id = $1 ORDER BY ingested_at DESC LIMIT $2")
            .bind(source_id).bind(limit).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn upsert(&self, entity: ExternalEntity) -> Result<ExternalEntity> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(entity.tenant_id).await?;
        let res = sqlx::query_as::<_, ExternalEntity>(
            r#"INSERT INTO external_entities (tenant_id, data_source_id, external_id, entity_type, raw_data, normalized_data, batch_id)
               VALUES ($1, $2, $3, $4, $5, $6, $7)
               ON CONFLICT (tenant_id, data_source_id, external_id) DO UPDATE SET
               entity_type = EXCLUDED.entity_type, raw_data = EXCLUDED.raw_data, normalized_data = EXCLUDED.normalized_data, batch_id = EXCLUDED.batch_id
               RETURNING *"#
        ).bind(entity.tenant_id).bind(entity.data_source_id).bind(entity.external_id).bind(entity.entity_type).bind(entity.raw_data).bind(entity.normalized_data).bind(entity.batch_id).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
}

#[async_trait]
pub trait CanonicalEntityRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<CanonicalEntity>>;
    async fn upsert(&self, entity: CanonicalEntity) -> Result<CanonicalEntity>;
    async fn resolve_external(&self, tenant_id: Uuid, system: &str, external_id: &str) -> Result<Option<CanonicalEntity>>;
    async fn invalidate_external(&self, tenant_id: Uuid, system: &str, external_id: &str) -> Result<()>;
    async fn set_locked(&self, id: Uuid, tenant_id: Uuid, locked: bool) -> Result<()>;
}

pub struct SqlxCanonicalEntityRepository { pool: PgPool }
impl SqlxCanonicalEntityRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl CanonicalEntityRepository for SqlxCanonicalEntityRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<CanonicalEntity>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, CanonicalEntity>("SELECT * FROM canonical_entities WHERE id = $1")
            .bind(id)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn upsert(&self, entity: CanonicalEntity) -> Result<CanonicalEntity> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(entity.tenant_id).await?;
        let res = sqlx::query_as::<_, CanonicalEntity>(
            r#"INSERT INTO canonical_entities (id, tenant_id, entity_type, canonical_data, field_provenance, confidence_score, updated_at, is_locked)
               VALUES ($1, $2, $3, $4, $5, $6, NOW(), $7)
               ON CONFLICT (id) DO UPDATE SET
               canonical_data = EXCLUDED.canonical_data, field_provenance = EXCLUDED.field_provenance, confidence_score = EXCLUDED.confidence_score, updated_at = NOW(), is_locked = EXCLUDED.is_locked
               RETURNING *"#
        ).bind(entity.id).bind(entity.tenant_id).bind(entity.entity_type).bind(entity.canonical_data).bind(entity.field_provenance).bind(entity.confidence_score).bind(entity.is_locked).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }

    async fn set_locked(&self, id: Uuid, tenant_id: Uuid, locked: bool) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("UPDATE canonical_entities SET is_locked = $2 WHERE id = $1")
            .bind(id)
            .bind(locked)
            .execute(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(())
    }
    async fn resolve_external(&self, tenant_id: Uuid, system: &str, external_id: &str) -> Result<Option<CanonicalEntity>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        // Optimized query to resolve from system name + external_id to canonical entity
        let res = sqlx::query_as::<_, CanonicalEntity>(
            r#"SELECT ce.* 
               FROM canonical_entities ce
               JOIN identity_links il ON ce.id = il.canonical_entity_id
               JOIN external_entities ee ON il.external_entity_id = ee.id
               JOIN data_sources ds ON ee.data_source_id = ds.id
               WHERE ds.name = $1 AND ee.external_id = $2"#
        )
        .bind(system)
        .bind(external_id)
        .fetch_optional(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }

    async fn invalidate_external(&self, _tenant_id: Uuid, _system: &str, _external_id: &str) -> Result<()> {
        // SQL implementation is stateless/no-cache
        Ok(())
    }
}

#[async_trait]
pub trait BatchRunRepository: Send + Sync {
    async fn create(&self, tenant_id: Uuid, job_type: &str) -> Result<BatchRun>;
    async fn update_status(&self, id: Uuid, status: &str, stats: serde_json::Value) -> Result<BatchRun>;
    async fn find_by_id(&self, id: Uuid) -> Result<Option<BatchRun>>;
    async fn list_by_tenant(&self, tenant_id: Uuid, limit: i64) -> Result<Vec<BatchRun>>;
}

pub struct CachedCanonicalEntityRepository<R: CanonicalEntityRepository> {
    inner: R,
    conn: redis::aio::MultiplexedConnection,
    local_cache: Option<moka::future::Cache<String, CanonicalEntity>>,
    ttl: u64,
}

impl<R: CanonicalEntityRepository> CachedCanonicalEntityRepository<R> {
    pub fn new(inner: R, conn: redis::aio::MultiplexedConnection) -> Self {
        Self {
            inner,
            conn,
            local_cache: None,
            ttl: 3600, // 1 hour default
        }
    }

    pub fn with_local_cache(mut self, cache: moka::future::Cache<String, CanonicalEntity>) -> Self {
        self.local_cache = Some(cache);
        self
    }

    fn cache_key(&self, tenant_id: Uuid, system: &str, external_id: &str) -> String {
        format!("v2:cannon:cache:resolve:{}:{}:{}", tenant_id, system, external_id)
    }
}

#[async_trait]
impl<R: CanonicalEntityRepository> CanonicalEntityRepository for CachedCanonicalEntityRepository<R> {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<CanonicalEntity>> {
        self.inner.find_by_id(id, tenant_id).await
    }

    async fn upsert(&self, entity: CanonicalEntity) -> Result<CanonicalEntity> {
        self.inner.upsert(entity).await
    }

    async fn resolve_external(&self, tenant_id: Uuid, system: &str, external_id: &str) -> Result<Option<CanonicalEntity>> {
        let key = self.cache_key(tenant_id, system, external_id);

        // 1. Try Local Cache (L1)
        if let Some(ref local) = self.local_cache {
            if let Some(entity) = local.get(&key).await {
                return Ok(Some(entity));
            }
        }

        let mut conn = self.conn.clone();

        // 2. Try Redis (L2) - Using bincode for binary speed
        if let Ok(cached) = conn.get::<&str, Vec<u8>>(&key).await {
            if let Ok(entity) = bincode::deserialize::<CanonicalEntity>(&cached) {
                // Populate Local Cache
                if let Some(ref local) = self.local_cache {
                    local.insert(key.clone(), entity.clone()).await;
                }
                return Ok(Some(entity));
            }
        }

        // 3. Fallback to Inner (DB)
        let result = self.inner.resolve_external(tenant_id, system, external_id).await?;

        // 4. Populate Caches
        if let Some(ref entity) = result {
            // Redis (L2)
            if let Ok(serialized) = bincode::serialize(entity) {
                let _: () = conn.set_ex::<_, _, ()>(&key, serialized, self.ttl).await.unwrap_or_default();
            }
            // Local Cache (L1)
            if let Some(ref local) = self.local_cache {
                local.insert(key, entity.clone()).await;
            }
        }

        Ok(result)
    }

    async fn invalidate_external(&self, tenant_id: Uuid, system: &str, external_id: &str) -> Result<()> {
        let key = self.cache_key(tenant_id, system, external_id);
        
        // Invalidate Local Cache
        if let Some(ref local) = self.local_cache {
            local.invalidate(&key).await;
        }

        // Invalidate Redis
        let mut conn = self.conn.clone();
        let _: () = conn.del(key).await.unwrap_or_default();
        
        self.inner.invalidate_external(tenant_id, system, external_id).await
    }

    async fn set_locked(&self, id: Uuid, tenant_id: Uuid, locked: bool) -> Result<()> {
        self.inner.set_locked(id, tenant_id, locked).await
    }
}

pub struct SqlxBatchRunRepository { pool: PgPool }
impl SqlxBatchRunRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }
#[async_trait]
pub trait IdentityLinkRepository: Send + Sync {
    async fn upsert(&self, link: IdentityLink) -> Result<IdentityLink>;
    async fn list_by_canonical(&self, canonical_id: Uuid, tenant_id: Uuid) -> Result<Vec<IdentityLink>>;
}

pub struct SqlxIdentityLinkRepository { pool: PgPool }
impl SqlxIdentityLinkRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl IdentityLinkRepository for SqlxIdentityLinkRepository {
    async fn upsert(&self, link: IdentityLink) -> Result<IdentityLink> {
        Ok(sqlx::query_as::<_, IdentityLink>(
            r#"INSERT INTO identity_links ( tenant_id, external_entity_id, canonical_entity_id, link_type, confidence, matched_on)
               VALUES ($1, $2, $3, $4, $5, $6)
               RETURNING *"#
        ).bind(link.tenant_id).bind(link.external_entity_id).bind(link.canonical_entity_id).bind(link.link_type).bind(link.confidence).bind(link.matched_on).fetch_one(&self.pool).await?)
    }
    async fn list_by_canonical(&self, canonical_id: Uuid, tenant_id: Uuid) -> Result<Vec<IdentityLink>> {
        Ok(sqlx::query_as::<_, IdentityLink>("SELECT * FROM identity_links WHERE canonical_entity_id = $1 AND tenant_id = $2").bind(canonical_id).bind(tenant_id).fetch_all(&self.pool).await?)
    }
}

pub struct InvalidatingIdentityLinkRepository<L: IdentityLinkRepository, C: CanonicalEntityRepository> {
    inner: L,
    canonical_repo: C,
    pool: PgPool,
}

impl<L: IdentityLinkRepository, C: CanonicalEntityRepository> InvalidatingIdentityLinkRepository<L, C> {
    pub fn new(inner: L, canonical_repo: C, pool: PgPool) -> Self {
        Self { inner, canonical_repo, pool }
    }
}

#[async_trait]
impl<L: IdentityLinkRepository, C: CanonicalEntityRepository> IdentityLinkRepository for InvalidatingIdentityLinkRepository<L, C> {
    async fn upsert(&self, link: IdentityLink) -> Result<IdentityLink> {
        let result = self.inner.upsert(link.clone()).await?;

        // Invalidate Cache for this external entity
        let ext_info: Option<(String, String)> = sqlx::query_as::<_, (String, String)>(
            r#"SELECT ds.name, ee.external_id 
               FROM external_entities ee
               JOIN data_sources ds ON ee.data_source_id = ds.id
               WHERE ee.id = $1 AND ee.tenant_id = $2"#
        )
        .bind(link.external_entity_id)
        .bind(link.tenant_id)
        .fetch_optional(&self.pool)
        .await?;

        if let Some((system, external_id)) = ext_info {
            let _ = self.canonical_repo.invalidate_external(link.tenant_id, &system, &external_id).await;
        }

        Ok(result)
    }

    async fn list_by_canonical(&self, canonical_id: Uuid, tenant_id: Uuid) -> Result<Vec<IdentityLink>> {
        self.inner.list_by_canonical(canonical_id, tenant_id).await
    }
}

#[async_trait]
pub trait MatchResultRepository: Send + Sync {
    async fn create(&self, result: MatchResult) -> Result<MatchResult>;
    async fn list_by_batch(&self, batch_id: Uuid, tenant_id: Uuid) -> Result<Vec<MatchResult>>;
}

pub struct SqlxMatchResultRepository { pool: PgPool }
impl SqlxMatchResultRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl MatchResultRepository for SqlxMatchResultRepository {
    async fn create(&self, result: MatchResult) -> Result<MatchResult> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(result.tenant_id).await?;
        let res = sqlx::query_as::<_, MatchResult>(
            r#"INSERT INTO match_results (tenant_id, batch_id, entity_a_id, entity_b_id, decision, confidence, rules_applied, explanation)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
               RETURNING *"#
        ).bind(result.tenant_id).bind(result.batch_id).bind(result.entity_a_id).bind(result.entity_b_id).bind(result.decision).bind(result.confidence).bind(result.rules_applied).bind(result.explanation).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_batch(&self, batch_id: Uuid, tenant_id: Uuid) -> Result<Vec<MatchResult>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, MatchResult>("SELECT * FROM match_results WHERE batch_id = $1").bind(batch_id).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
}

#[async_trait]
pub trait MatchRuleRepository: Send + Sync {
    async fn find_active_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<MatchRule>>;
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<MatchRule>>;
    async fn list_versions_by_name(&self, name: &str, tenant_id: Uuid) -> Result<Vec<MatchRule>>;
    async fn create_version(&self, rule: MatchRule) -> Result<MatchRule>;
    async fn deactivate_others_by_name(&self, name: &str, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxMatchRuleRepository { pool: PgPool }
impl SqlxMatchRuleRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl MatchRuleRepository for SqlxMatchRuleRepository {
    async fn find_active_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<MatchRule>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, MatchRule>("SELECT * FROM match_rules WHERE is_active = TRUE").fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<MatchRule>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, MatchRule>("SELECT * FROM match_rules WHERE id = $1").bind(id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_versions_by_name(&self, name: &str, tenant_id: Uuid) -> Result<Vec<MatchRule>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, MatchRule>("SELECT * FROM match_rules WHERE name = $1 ORDER BY version DESC").bind(name).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create_version(&self, rule: MatchRule) -> Result<MatchRule> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(rule.tenant_id).await?;
        let next_version = sqlx::query_scalar::<_, i32>("SELECT COALESCE(MAX(version), 0) + 1 FROM match_rules WHERE name = $1")
            .bind(&rule.name).fetch_one(&mut *tx).await?;
        
        let res = sqlx::query_as::<_, MatchRule>(
            r#"INSERT INTO match_rules (id, tenant_id, version, name, rule_type, config, weight, is_active, created_by)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
               RETURNING *"#
        ).bind(Uuid::new_v4()).bind(rule.tenant_id).bind(next_version).bind(rule.name).bind(rule.rule_type).bind(rule.config).bind(rule.weight).bind(rule.is_active).bind(rule.created_by).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn deactivate_others_by_name(&self, name: &str, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("UPDATE match_rules SET is_active = FALSE WHERE name = $1").bind(name).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait UserRepository: Send + Sync {
    async fn find_by_email(&self, email: &str, tenant_id: Uuid) -> Result<Option<User>>;
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<User>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<User>>;
    async fn create(&self, user: User) -> Result<User>;
    async fn update_role(&self, id: Uuid, tenant_id: Uuid, role: &str) -> Result<User>;
    async fn update_password(&self, id: Uuid, tenant_id: Uuid, hashed_password: &str) -> Result<User>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxUserRepository { pool: PgPool }
impl SqlxUserRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl UserRepository for SqlxUserRepository {
    async fn find_by_email(&self, email: &str, tenant_id: Uuid) -> Result<Option<User>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, User>("SELECT * FROM users WHERE email = $1").bind(email).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<User>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = $1").bind(id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<User>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, User>("SELECT * FROM users ORDER BY created_at DESC").fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create(&self, user: User) -> Result<User> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(user.tenant_id).await?;
        let res = sqlx::query_as::<_, User>(
            r#"INSERT INTO users (tenant_id, email, name, role, hashed_password, external_id, provisioning_source, must_change_password, is_verified, verification_token)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
               RETURNING *"#
        )
        .bind(user.tenant_id)
        .bind(user.email)
        .bind(user.name)
        .bind(user.role)
        .bind(user.hashed_password)
        .bind(user.external_id)
        .bind(user.provisioning_source)
        .bind(user.must_change_password)
        .bind(user.is_verified)
        .bind(user.verification_token)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }

    async fn update_role(&self, id: Uuid, tenant_id: Uuid, role: &str) -> Result<User> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, User>("UPDATE users SET role = $2 WHERE id = $1 RETURNING *").bind(id).bind(role).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update_password(&self, id: Uuid, tenant_id: Uuid, hashed_password: &str) -> Result<User> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, User>("UPDATE users SET hashed_password = $2, must_change_password = FALSE WHERE id = $1 RETURNING *").bind(id).bind(hashed_password).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM users WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait MatchAuditRepository: Send + Sync {
    async fn create(&self, audit: MatchAudit) -> Result<()>;
    async fn find_by_batch(&self, tenant_id: Uuid, batch_id: Uuid) -> Result<Vec<MatchAudit>>;
    async fn find_by_entity(&self, tenant_id: Uuid, entity_id: Uuid) -> Result<Vec<MatchAudit>>;
}

pub struct SqlxMatchAuditRepository { pool: PgPool }
impl SqlxMatchAuditRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl MatchAuditRepository for SqlxMatchAuditRepository {
    async fn create(&self, audit: MatchAudit) -> Result<()> {
        sqlx::query(
            r#"INSERT INTO match_audit (id, tenant_id, batch_run_id, entity_a_id, entity_b_id, decision, confidence, rule_trace, temporal_adjustment, created_at)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)"#
        )
        .bind(audit.id)
        .bind(audit.tenant_id)
        .bind(audit.batch_run_id)
        .bind(audit.entity_a_id)
        .bind(audit.entity_b_id)
        .bind(audit.decision)
        .bind(audit.confidence)
        .bind(audit.rule_trace)
        .bind(audit.temporal_adjustment)
        .bind(audit.created_at)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    async fn find_by_batch(&self, tenant_id: Uuid, batch_id: Uuid) -> Result<Vec<MatchAudit>> {
        let rows = sqlx::query_as::<_, MatchAudit>(
            "SELECT * FROM match_audit WHERE tenant_id = $1 AND batch_run_id = $2 ORDER BY created_at DESC"
        )
        .bind(tenant_id)
        .bind(batch_id)
        .fetch_all(&self.pool)
        .await?;
        Ok(rows)
    }

    async fn find_by_entity(&self, tenant_id: Uuid, entity_id: Uuid) -> Result<Vec<MatchAudit>> {
        let rows = sqlx::query_as::<_, MatchAudit>(
            "SELECT * FROM match_audit WHERE tenant_id = $1 AND (entity_a_id = $2 OR entity_b_id = $2) ORDER BY created_at DESC"
        )
        .bind(tenant_id)
        .bind(entity_id)
        .fetch_all(&self.pool)
        .await?;
        Ok(rows)
    }
}

#[async_trait]
pub trait ManualOverrideRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<ManualOverride>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<ManualOverride>>;
    async fn create(&self, override_item: ManualOverride) -> Result<ManualOverride>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxManualOverrideRepository { pool: PgPool }
impl SqlxManualOverrideRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl ManualOverrideRepository for SqlxManualOverrideRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<ManualOverride>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ManualOverride>("SELECT * FROM manual_overrides WHERE id = $1").bind(id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<ManualOverride>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ManualOverride>("SELECT * FROM manual_overrides WHERE superseded_at IS NULL").fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create(&self, o: ManualOverride) -> Result<ManualOverride> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(o.tenant_id).await?;
        let res = sqlx::query_as::<_, ManualOverride>(
            r#"INSERT INTO manual_overrides (tenant_id, canonical_entity_id, override_type, override_data, reason, created_by)
               VALUES ($1, $2, $3, $4, $5, $6)
               RETURNING *"#
        ).bind(o.tenant_id).bind(o.canonical_entity_id).bind(o.override_type).bind(o.override_data).bind(o.reason).bind(o.created_by).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("UPDATE manual_overrides SET superseded_at = NOW() WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
impl BatchRunRepository for SqlxBatchRunRepository {
    async fn create(&self, tenant_id: Uuid, job_type: &str) -> Result<BatchRun> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, BatchRun>("INSERT INTO batch_runs (tenant_id, status, job_type, started_at) VALUES ($1, 'running', $2, NOW()) RETURNING *").bind(tenant_id).bind(job_type).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update_status(&self, id: Uuid, status: &str, stats: serde_json::Value) -> Result<BatchRun> {
        let batch = sqlx::query_as::<_, BatchRun>("SELECT * FROM batch_runs WHERE id = $1").bind(id).fetch_one(&self.pool).await?;
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(batch.tenant_id).await?;
        let res = sqlx::query_as::<_, BatchRun>("UPDATE batch_runs SET status = $2, stats = $3, completed_at = CASE WHEN $2 IN ('completed', 'failed') THEN NOW() ELSE completed_at END WHERE id = $1 RETURNING *").bind(id).bind(status).bind(stats).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn find_by_id(&self, id: Uuid) -> Result<Option<BatchRun>> {
        Ok(sqlx::query_as::<_, BatchRun>("SELECT * FROM batch_runs WHERE id = $1").bind(id).fetch_optional(&self.pool).await?)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid, limit: i64) -> Result<Vec<BatchRun>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, BatchRun>("SELECT * FROM batch_runs ORDER BY created_at DESC LIMIT $1").bind(limit).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
}

#[async_trait]
pub trait AuditEventRepository: Send + Sync {
    async fn create(&self, event: AuditEvent) -> Result<AuditEvent>;
    async fn list_recent(&self, tenant_id: Uuid, limit: i64) -> Result<Vec<AuditEvent>>;
    async fn list_by_resource(&self, tenant_id: Uuid, resource_id: Uuid) -> Result<Vec<AuditEvent>>;
}

pub struct SqlxAuditEventRepository { 
    pool: PgPool,
    redis: Option<redis::Client>,
}

impl SqlxAuditEventRepository { 
    pub fn new(pool: PgPool) -> Self { Self { pool, redis: None } }
    pub fn with_redis(pool: PgPool, redis: redis::Client) -> Self { Self { pool, redis: Some(redis) } }
}

#[async_trait]
impl AuditEventRepository for SqlxAuditEventRepository {
    async fn create(&self, event: AuditEvent) -> Result<AuditEvent> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(event.tenant_id).await?;
        let created = sqlx::query_as::<_, AuditEvent>(
            r#"INSERT INTO audit_events (tenant_id, actor_id, actor_type, action, resource_type, resource_id, before_state, after_state, reason, timestamp)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
               RETURNING *"#
        ).bind(event.tenant_id).bind(event.actor_id).bind(event.actor_type).bind(event.action).bind(event.resource_type).bind(event.resource_id).bind(event.before_state).bind(event.after_state).bind(event.reason).bind(event.timestamp).fetch_one(&mut *tx).await?;
        tx.commit().await?;

        if let Some(redis) = &self.redis {
            use redis::AsyncCommands;
            if let Ok(mut conn) = redis.get_multiplexed_async_connection().await {
                let channel = format!("audit_events:{}", created.tenant_id);
                let payload = serde_json::to_string(&created).unwrap_or_default();
                let _: std::result::Result<(), _> = conn.publish(channel, payload).await;
            }
        }

        Ok(created)
    }
    async fn list_recent(&self, tenant_id: Uuid, limit: i64) -> Result<Vec<AuditEvent>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, AuditEvent>("SELECT * FROM audit_events WHERE tenant_id = $1 ORDER BY timestamp DESC LIMIT $2").bind(tenant_id).bind(limit).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_resource(&self, tenant_id: Uuid, resource_id: Uuid) -> Result<Vec<AuditEvent>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, AuditEvent>("SELECT * FROM audit_events WHERE tenant_id = $1 AND resource_id = $2 ORDER BY timestamp DESC").bind(tenant_id).bind(resource_id).fetch_all(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
}

#[async_trait]
#[async_trait]
pub trait RefreshTokenRepository: Send + Sync {
    async fn create(&self, user_id: Uuid, tenant_id: Uuid, ttl_days: u32, family_id: Option<Uuid>, parent_id: Option<Uuid>) -> Result<(String, Uuid)>;
    async fn find_by_hash(&self, token_hash: &str) -> Result<Option<RefreshToken>>;
    async fn update_replaced(&self, id: Uuid) -> Result<()>;
    async fn revoke_family(&self, family_id: Uuid) -> Result<()>;
    async fn delete_family(&self, family_id: Uuid) -> Result<()>;
    async fn cleanup(&self, retention_days: u32) -> Result<u64>;
}

pub struct SqlxRefreshTokenRepository { pool: PgPool }
impl SqlxRefreshTokenRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl RefreshTokenRepository for SqlxRefreshTokenRepository {
    async fn create(&self, user_id: Uuid, tenant_id: Uuid, ttl_days: u32, family_id: Option<Uuid>, parent_id: Option<Uuid>) -> Result<(String, Uuid)> {
        // Use ring for CSPRNG
        use ring::rand::{SystemRandom, SecureRandom};
        let rng = SystemRandom::new();
        let mut random_bytes = [0u8; 32];
        rng.fill(&mut random_bytes).map_err(|_| crate::error::CannonError::Internal("RNG failure".into()))?;
        let plaintext_token = hex::encode(random_bytes);

        let token_hash = sha2::Sha256::digest(plaintext_token.as_bytes());
        let hash_hex = hex::encode(token_hash);
        
        let actual_family_id = family_id.unwrap_or_else(Uuid::new_v4);

        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query(
            "INSERT INTO refresh_tokens (token_hash, user_id, tenant_id, family_id, parent_id, expires_at) VALUES ($1, $2, $3, $4, $5, NOW() + ($6 || ' days')::interval)"
        )
        .bind(&hash_hex)
        .bind(user_id)
        .bind(tenant_id)
        .bind(actual_family_id)
        .bind(parent_id)
        .bind(ttl_days as i32)
        .execute(&mut *tx)
        .await?;
        tx.commit().await?;
        
        Ok((plaintext_token, actual_family_id))
    }

    async fn find_by_hash(&self, token_hash: &str) -> Result<Option<RefreshToken>> {
        // We find it even if revoked, to detect reuse
        Ok(sqlx::query_as::<_, RefreshToken>("SELECT * FROM refresh_tokens WHERE token_hash = $1")
            .bind(token_hash)
            .fetch_optional(&self.pool)
            .await?)
    }

    async fn update_replaced(&self, id: Uuid) -> Result<()> {
        let result = sqlx::query("UPDATE refresh_tokens SET replaced_at = NOW() WHERE id = $1 AND replaced_at IS NULL")
            .bind(id)
            .execute(&self.pool)
            .await?;
            
        if result.rows_affected() == 0 {
             return Err(crate::error::CannonError::Unauthorized.into()); // Treat as conflict/theft
        }
        Ok(())
    }

    async fn revoke_family(&self, family_id: Uuid) -> Result<()> {
        sqlx::query("UPDATE refresh_tokens SET revoked_at = NOW() WHERE family_id = $1")
            .bind(family_id)
            .execute(&self.pool)
            .await?;
        Ok(())
    }

    async fn delete_family(&self, family_id: Uuid) -> Result<()> {
        sqlx::query("DELETE FROM refresh_tokens WHERE family_id = $1")
            .bind(family_id)
            .execute(&self.pool)
            .await?;
        Ok(())
    }

    async fn cleanup(&self, retention_days: u32) -> Result<u64> {
        let result = sqlx::query("DELETE FROM refresh_tokens WHERE (revoked_at IS NOT NULL OR replaced_at IS NOT NULL) AND created_at < NOW() - ($1 || ' days')::interval")
            .bind(retention_days as i32)
            .execute(&self.pool)
            .await?;
        Ok(result.rows_affected())
    }
}

#[async_trait]
pub trait SsoConfigRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<SsoConfig>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<SsoConfig>>;
    async fn create(&self, config: SsoConfig) -> Result<SsoConfig>;
    async fn update(&self, config: SsoConfig) -> Result<SsoConfig>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxSsoConfigRepository { pool: PgPool }
impl SqlxSsoConfigRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl SsoConfigRepository for SqlxSsoConfigRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<SsoConfig>> {
        Ok(sqlx::query_as::<_, SsoConfig>("SELECT * FROM sso_configurations WHERE id = $1 AND tenant_id = $2")
            .bind(id)
            .bind(tenant_id)
            .fetch_optional(&self.pool)
            .await?)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<SsoConfig>> {
        Ok(sqlx::query_as::<_, SsoConfig>("SELECT * FROM sso_configurations WHERE tenant_id = $1")
            .bind(tenant_id)
            .fetch_all(&self.pool)
            .await?)
    }
    async fn create(&self, config: SsoConfig) -> Result<SsoConfig> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(config.tenant_id).await?;
        let res = sqlx::query_as::<_, SsoConfig>(
            r#"INSERT INTO sso_configurations (tenant_id, provider_type, display_name, issuer_url, client_id, client_secret_enc, auth_url, token_url, jwks_url, saml_metadata_url, is_active)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
               RETURNING *"#
        )
        .bind(config.tenant_id)
        .bind(config.provider_type)
        .bind(config.display_name)
        .bind(config.issuer_url)
        .bind(config.client_id)
        .bind(config.client_secret_enc)
        .bind(config.auth_url)
        .bind(config.token_url)
        .bind(config.jwks_url)
        .bind(config.saml_metadata_url)
        .bind(config.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update(&self, config: SsoConfig) -> Result<SsoConfig> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(config.tenant_id).await?;
        let res = sqlx::query_as::<_, SsoConfig>(
            r#"UPDATE sso_configurations 
               SET provider_type = $3, display_name = $4, issuer_url = $5, client_id = $6, client_secret_enc = $7, auth_url = $8, token_url = $9, jwks_url = $10, saml_metadata_url = $11, is_active = $12, updated_at = NOW()
               WHERE id = $1
               RETURNING *"#
        )
        .bind(config.id)
        .bind(config.tenant_id)
        .bind(config.provider_type)
        .bind(config.display_name)
        .bind(config.issuer_url)
        .bind(config.client_id)
        .bind(config.client_secret_enc)
        .bind(config.auth_url)
        .bind(config.token_url)
        .bind(config.jwks_url)
        .bind(config.saml_metadata_url)
        .bind(config.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM sso_configurations WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait AuditStreamRepository: Send + Sync {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<AuditStream>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<AuditStream>>;
    async fn create(&self, stream: AuditStream) -> Result<AuditStream>;
    async fn update(&self, stream: AuditStream) -> Result<AuditStream>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxAuditStreamRepository { pool: PgPool }
impl SqlxAuditStreamRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl AuditStreamRepository for SqlxAuditStreamRepository {
    async fn find_by_id(&self, id: Uuid, tenant_id: Uuid) -> Result<Option<AuditStream>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, AuditStream>("SELECT * FROM audit_streams WHERE id = $1")
            .bind(id)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<AuditStream>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, AuditStream>("SELECT * FROM audit_streams")
            .fetch_all(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create(&self, stream: AuditStream) -> Result<AuditStream> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(stream.tenant_id).await?;
        let res = sqlx::query_as::<_, AuditStream>(
            r#"INSERT INTO audit_streams (tenant_id, destination_type, destination_url, auth_config, is_active)
               VALUES ($1, $2, $3, $4, $5)
               RETURNING *"#
        )
        .bind(stream.tenant_id)
        .bind(stream.destination_type)
        .bind(stream.destination_url)
        .bind(stream.auth_config)
        .bind(stream.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn update(&self, stream: AuditStream) -> Result<AuditStream> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(stream.tenant_id).await?;
        let res = sqlx::query_as::<_, AuditStream>(
            r#"UPDATE audit_streams 
               SET destination_type = $3, destination_url = $4, auth_config = $5, is_active = $6, updated_at = NOW()
               WHERE id = $1
               RETURNING *"#
        )
        .bind(stream.id)
        .bind(stream.tenant_id)
        .bind(stream.destination_type)
        .bind(stream.destination_url)
        .bind(stream.auth_config)
        .bind(stream.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM audit_streams WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait RetentionPolicyRepository: Send + Sync {
    async fn find_by_entity_type(&self, tenant_id: Uuid, entity_type: &str) -> Result<Option<RetentionPolicy>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<RetentionPolicy>>;
    async fn upsert(&self, policy: RetentionPolicy) -> Result<RetentionPolicy>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxRetentionPolicyRepository { pool: PgPool }
impl SqlxRetentionPolicyRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl RetentionPolicyRepository for SqlxRetentionPolicyRepository {
    async fn find_by_entity_type(&self, tenant_id: Uuid, entity_type: &str) -> Result<Option<RetentionPolicy>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, RetentionPolicy>("SELECT * FROM retention_policies WHERE entity_type = $1")
            .bind(entity_type)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<RetentionPolicy>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, RetentionPolicy>("SELECT * FROM retention_policies")
            .fetch_all(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn upsert(&self, p: RetentionPolicy) -> Result<RetentionPolicy> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(p.tenant_id).await?;
        let res = sqlx::query_as::<_, RetentionPolicy>(
            r#"INSERT INTO retention_policies (tenant_id, entity_type, retention_days, shred_pii, is_active)
               VALUES ($1, $2, $3, $4, $5)
               ON CONFLICT (tenant_id, entity_type) DO UPDATE SET
               retention_days = EXCLUDED.retention_days, shred_pii = EXCLUDED.shred_pii, is_active = EXCLUDED.is_active, updated_at = NOW()
               RETURNING *"#
        )
        .bind(p.tenant_id)
        .bind(p.entity_type)
        .bind(p.retention_days)
        .bind(p.shred_pii)
        .bind(p.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM retention_policies WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait ApiKeyRepository: Send + Sync {
    async fn find_by_hash(&self, hash: &str) -> Result<Option<ApiKey>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<ApiKey>>;
    async fn create(&self, key: ApiKey) -> Result<ApiKey>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
    async fn update_last_used(&self, id: Uuid) -> Result<()>;
}

pub struct SqlxApiKeyRepository { pool: PgPool }
impl SqlxApiKeyRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl ApiKeyRepository for SqlxApiKeyRepository {
    async fn find_by_hash(&self, hash: &str) -> Result<Option<ApiKey>> {
        // find_by_hash is used for auth, might not have tenant_id.
        // It bypassess RLS or uses its own logic.
        Ok(sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys WHERE key_hash = $1")
            .bind(hash)
            .fetch_optional(&self.pool)
            .await?)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<ApiKey>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, ApiKey>("SELECT * FROM api_keys ORDER BY created_at DESC")
            .fetch_all(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn create(&self, k: ApiKey) -> Result<ApiKey> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(k.tenant_id).await?;
        let res = sqlx::query_as::<_, ApiKey>(
            r#"INSERT INTO api_keys (id, tenant_id, name, key_hash, key_prefix, scopes, expires_at)
               VALUES ($1, $2, $3, $4, $5, $6, $7)
               RETURNING *"#
        )
        .bind(k.id)
        .bind(k.tenant_id)
        .bind(k.name)
        .bind(k.key_hash)
        .bind(k.key_prefix)
        .bind(k.scopes)
        .bind(k.expires_at)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM api_keys WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
    async fn update_last_used(&self, id: Uuid) -> Result<()> {
        // update_last_used doesn't have tenant_id, might need workaround.
        // For now, allow it to update any key.
        sqlx::query("UPDATE api_keys SET last_used_at = NOW() WHERE id = $1")
            .bind(id)
            .execute(&self.pool)
            .await?;
        Ok(())
    }
}

#[async_trait]
pub trait SurvivorshipPolicyRepository: Send + Sync {
    async fn find_by_field(&self, tenant_id: Uuid, field_name: &str) -> Result<Option<SurvivorshipPolicy>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<SurvivorshipPolicy>>;
    async fn upsert(&self, policy: SurvivorshipPolicy) -> Result<SurvivorshipPolicy>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxSurvivorshipPolicyRepository { pool: PgPool }
impl SqlxSurvivorshipPolicyRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl SurvivorshipPolicyRepository for SqlxSurvivorshipPolicyRepository {
    async fn find_by_field(&self, tenant_id: Uuid, field_name: &str) -> Result<Option<SurvivorshipPolicy>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, SurvivorshipPolicy>("SELECT * FROM survivorship_policies WHERE field_name = $1")
            .bind(field_name)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<SurvivorshipPolicy>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, SurvivorshipPolicy>("SELECT * FROM survivorship_policies")
            .fetch_all(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn upsert(&self, p: SurvivorshipPolicy) -> Result<SurvivorshipPolicy> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(p.tenant_id).await?;
        let res = sqlx::query_as::<_, SurvivorshipPolicy>(
            r#"INSERT INTO survivorship_policies (id, tenant_id, field_name, priority_config, is_active)
               VALUES ($1, $2, $3, $4, $5)
               ON CONFLICT (tenant_id, field_name) DO UPDATE SET
               priority_config = EXCLUDED.priority_config, is_active = EXCLUDED.is_active, updated_at = NOW()
               RETURNING *"#
        )
        .bind(p.id)
        .bind(p.tenant_id)
        .bind(p.field_name)
        .bind(p.priority_config)
        .bind(p.is_active)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM survivorship_policies WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait SchemaMappingRepository: Send + Sync {
    async fn find_by_source(&self, tenant_id: Uuid, data_source_id: Uuid) -> Result<Option<SchemaMapping>>;
    async fn upsert(&self, mapping: SchemaMapping) -> Result<SchemaMapping>;
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()>;
}

pub struct SqlxSchemaMappingRepository { pool: PgPool }
impl SqlxSchemaMappingRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl SchemaMappingRepository for SqlxSchemaMappingRepository {
    async fn find_by_source(&self, tenant_id: Uuid, data_source_id: Uuid) -> Result<Option<SchemaMapping>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, SchemaMapping>("SELECT * FROM schema_mappings WHERE data_source_id = $1")
            .bind(data_source_id)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn upsert(&self, m: SchemaMapping) -> Result<SchemaMapping> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(m.tenant_id).await?;
        let res = sqlx::query_as::<_, SchemaMapping>(
            r#"INSERT INTO schema_mappings (id, tenant_id, data_source_id, mapping_config)
               VALUES ($1, $2, $3, $4)
               ON CONFLICT (tenant_id, data_source_id) DO UPDATE SET
               mapping_config = EXCLUDED.mapping_config, updated_at = NOW()
               RETURNING *"#
        )
        .bind(m.id)
        .bind(m.tenant_id)
        .bind(m.data_source_id)
        .bind(m.mapping_config)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn delete(&self, id: Uuid, tenant_id: Uuid) -> Result<()> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        sqlx::query("DELETE FROM schema_mappings WHERE id = $1").bind(id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
}

#[async_trait]
pub trait IdentityPlanRepository: Send + Sync {
    async fn upsert(&self, tenant_id: Uuid, version: &str, plan_hash: &str, raw_yaml: &str, compiled_plan: serde_json::Value) -> Result<IdentityPlanModel>;
    async fn find_by_version(&self, tenant_id: Uuid, version: &str) -> Result<Option<IdentityPlanModel>>;
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<IdentityPlanModel>>;
}

pub struct SqlxIdentityPlanRepository { pool: PgPool }
impl SqlxIdentityPlanRepository { pub fn new(pool: PgPool) -> Self { Self { pool } } }

#[async_trait]
impl IdentityPlanRepository for SqlxIdentityPlanRepository {
    async fn upsert(&self, tenant_id: Uuid, version: &str, plan_hash: &str, raw_yaml: &str, compiled_plan: serde_json::Value) -> Result<IdentityPlanModel> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, IdentityPlanModel>(
            r#"INSERT INTO identity_plans (tenant_id, identity_version, plan_hash, raw_yaml, compiled_plan, created_at)
               VALUES ($1, $2, $3, $4, $5, NOW())
               ON CONFLICT (tenant_id, identity_version) DO UPDATE SET
               plan_hash = EXCLUDED.plan_hash, raw_yaml = EXCLUDED.raw_yaml, compiled_plan = EXCLUDED.compiled_plan, created_at = NOW()
               RETURNING *"#
        )
        .bind(tenant_id)
        .bind(version)
        .bind(plan_hash)
        .bind(raw_yaml)
        .bind(compiled_plan)
        .fetch_one(&mut *tx)
        .await?;
        tx.commit().await?;
        Ok(res)
    }
    
    async fn find_by_version(&self, tenant_id: Uuid, version: &str) -> Result<Option<IdentityPlanModel>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, IdentityPlanModel>("SELECT * FROM identity_plans WHERE identity_version = $1")
            .bind(version)
            .fetch_optional(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
    
    async fn list_by_tenant(&self, tenant_id: Uuid) -> Result<Vec<IdentityPlanModel>> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(tenant_id).await?;
        let res = sqlx::query_as::<_, IdentityPlanModel>("SELECT * FROM identity_plans ORDER BY created_at DESC")
            .fetch_all(&mut *tx)
            .await?;
        tx.commit().await?;
        Ok(res)
    }
}
