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<()> {
        let tid = tenant_id.to_string();
        sqlx::query!("SELECT set_config('request.tenant_id', $1::text, true)", tid)
            .fetch_one(&mut *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<()> {
        let tid = tenant_id.to_string();
        sqlx::query!("SELECT set_config('request.tenant_id', $1::text, true)", tid)
            .fetch_one(&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?;
        let res = sqlx::query_as!(Tenant, "SELECT * FROM tenants WHERE id = $1", id).fetch_optional(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }
    async fn find_by_slug(&self, slug: &str) -> Result<Option<Tenant>> {
        Ok(sqlx::query_as!(Tenant, "SELECT * FROM tenants WHERE slug = $1", 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 *", name, slug, parent_id, tier, 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", 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 *", id, tier, 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 *", id, 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 *",
            tenant.id, tenant.name, tenant.slug, tenant.settings, tenant.tier, tenant.usage_limits, 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", 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 *", tenant_id, name, source_type, 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 *", id, name, config, 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)))
    }
    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", id, 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)));
        }
        Ok(())
    }
}

/// Status returned by idempotent upsert: whether the record was new, updated, or unchanged.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UpsertStatus {
    New,
    Updated,
    Unchanged,
}

#[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>;
    /// Idempotent upsert: skips write if content_hash matches existing record.
    async fn upsert_idempotent(&self, entity: ExternalEntity) -> Result<(ExternalEntity, UpsertStatus)>;
}

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", 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", entity_type, limit, 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", source_id, 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, content_hash)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
               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, content_hash = EXCLUDED.content_hash
               RETURNING *"#,
            entity.tenant_id, entity.data_source_id, entity.external_id, entity.entity_type, entity.raw_data, entity.normalized_data, entity.batch_id, entity.content_hash
        ).fetch_one(&mut *tx).await?;
        tx.commit().await?;
        Ok(res)
    }

    async fn upsert_idempotent(&self, entity: ExternalEntity) -> Result<(ExternalEntity, UpsertStatus)> {
        let mut tx = self.pool.begin().await?;
        tx.set_tenant_context(entity.tenant_id).await?;

        let existing: Option<ExternalEntity> = sqlx::query_as!(ExternalEntity,
            "SELECT * FROM external_entities WHERE tenant_id = $1 AND data_source_id = $2 AND external_id = $3",
            entity.tenant_id, entity.data_source_id, &entity.external_id
        )
        .fetch_optional(&mut *tx)
        .await?;

        match existing {
            Some(ref row)
                if row.content_hash.is_some()
                    && entity.content_hash == row.content_hash =>
            {
                let ent = row.clone();
                tx.commit().await?;
                Ok((ent, UpsertStatus::Unchanged))
            }
            Some(_) => {
                let res = sqlx::query_as!(ExternalEntity,
                    r#"UPDATE external_entities SET
                       entity_type = $4, raw_data = $5, normalized_data = $6, batch_id = $7, content_hash = $8, ingested_at = NOW()
                       WHERE tenant_id = $1 AND data_source_id = $2 AND external_id = $3
                       RETURNING *"#,
                    entity.tenant_id, entity.data_source_id, &entity.external_id,
                    &entity.entity_type, &entity.raw_data, entity.normalized_data, entity.batch_id, entity.content_hash.as_deref()
                )
                .fetch_one(&mut *tx)
                .await?;
                tx.commit().await?;
                Ok((res, UpsertStatus::Updated))
            }
            None => {
                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, content_hash)
                       VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                       RETURNING *"#,
                    entity.tenant_id, entity.data_source_id, &entity.external_id,
                    &entity.entity_type, &entity.raw_data, entity.normalized_data, entity.batch_id, entity.content_hash.as_deref()
                )
                .fetch_one(&mut *tx)
                .await?;
                tx.commit().await?;
                Ok((res, UpsertStatus::New))
            }
        }
    }
}

#[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", 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, identity_version)
               VALUES ($1, $2, $3, $4, $5, $6, NOW(), $7, $8)
               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,
               identity_version = COALESCE(EXCLUDED.identity_version, canonical_entities.identity_version)
               RETURNING *"#,
            entity.id, entity.tenant_id, entity.entity_type, entity.canonical_data, entity.field_provenance, entity.confidence_score, entity.is_locked, entity.identity_version
        ).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", id, 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?;
        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"#,
            system, 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<()> {
        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,
        }
    }

    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);

        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();

        if let Ok(cached) = conn.get::<&str, Vec<u8>>(&key).await {
            if let Ok(entity) = bincode::deserialize::<CanonicalEntity>(&cached) {
                if let Some(ref local) = self.local_cache {
                    local.insert(key.clone(), entity.clone()).await;
                }
                return Ok(Some(entity));
            }
        }

        let result = self.inner.resolve_external(tenant_id, system, external_id).await?;

        if let Some(ref entity) = result {
            if let Ok(serialized) = bincode::serialize(entity) {
                let _: () = conn.set_ex::<_, _, ()>(&key, serialized, self.ttl).await.unwrap_or_default();
            }
            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);

        if let Some(ref local) = self.local_cache {
            local.invalidate(&key).await;
        }

        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 *"#,
            link.tenant_id, link.external_entity_id, link.canonical_entity_id, link.link_type, link.confidence, 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", canonical_id, 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?;

        // Tuple queries can't use query_as! macro - use query! with named fields
        let ext_info = sqlx::query!(
            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"#,
            link.external_entity_id, link.tenant_id
        )
        .fetch_optional(&self.pool)
        .await?;

        if let Some(row) = ext_info {
            let _ = self.canonical_repo.invalidate_external(link.tenant_id, &row.name, &row.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 *"#,
            result.tenant_id, result.batch_id, result.entity_a_id, result.entity_b_id, result.decision, result.confidence, result.rules_applied, 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", 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", 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", 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: Option<i32> = sqlx::query_scalar!("SELECT COALESCE(MAX(version), 0) + 1 FROM match_rules WHERE name = $1", &rule.name).fetch_one(&mut *tx).await?;
        let next_version = next_version.unwrap_or(1);
        let new_id = Uuid::new_v4();
        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 *"#,
            new_id, rule.tenant_id, next_version, rule.name, rule.rule_type, rule.config, rule.weight, rule.is_active, 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", 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", 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", 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 *"#,
            user.tenant_id, user.email, user.name, user.role, user.hashed_password,
            user.external_id, user.provisioning_source, user.must_change_password, user.is_verified, 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 *", id, 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 *", id, 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", 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)"#,
            audit.id, audit.tenant_id, audit.batch_run_id, audit.entity_a_id, audit.entity_b_id,
            audit.decision, audit.confidence, audit.rule_trace, audit.temporal_adjustment, 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 id, tenant_id, batch_run_id, entity_a_id, entity_b_id, decision, confidence, rule_trace, temporal_adjustment, created_at FROM match_audit WHERE tenant_id = $1 AND batch_run_id = $2 ORDER BY created_at DESC",
            tenant_id, 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 id, tenant_id, batch_run_id, entity_a_id, entity_b_id, decision, confidence, rule_trace, temporal_adjustment, created_at FROM match_audit WHERE tenant_id = $1 AND (entity_a_id = $2 OR entity_b_id = $2) ORDER BY created_at DESC",
            tenant_id, 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", 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 *"#,
            o.tenant_id, o.canonical_entity_id, o.override_type, o.override_data, o.reason, 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", 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 *", tenant_id, 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", 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::text, stats = $3, completed_at = CASE WHEN $2::text IN ('completed', 'failed') THEN NOW() ELSE completed_at END WHERE id = $1 RETURNING *", id, status, 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", 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", 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 *"#,
            event.tenant_id, event.actor_id, event.actor_type, event.action, event.resource_type, event.resource_id, event.before_state, event.after_state, event.reason, 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", tenant_id, 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", tenant_id, 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::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 ttl_days_i32 = ttl_days as i32;

        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)",
            hash_hex, user_id, tenant_id, actual_family_id, parent_id, ttl_days_i32.to_string()
        )
        .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>> {
        Ok(sqlx::query_as!(RefreshToken, "SELECT * FROM refresh_tokens WHERE token_hash = $1", 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", id)
            .execute(&self.pool)
            .await?;

        if result.rows_affected() == 0 {
             return Err(crate::error::CannonError::Unauthorized);
        }
        Ok(())
    }

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

    async fn cleanup(&self, retention_days: u32) -> Result<u64> {
        let days_str = (retention_days as i32).to_string();
        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", days_str)
            .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", id, 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", 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 *"#,
            config.tenant_id, config.provider_type, config.display_name, config.issuer_url, config.client_id, config.client_secret_enc, config.auth_url, config.token_url, config.jwks_url, config.saml_metadata_url, 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 AND tenant_id = $2
               RETURNING *"#,
            config.id, config.tenant_id, config.provider_type, config.display_name, config.issuer_url, config.client_id, config.client_secret_enc, config.auth_url, config.token_url, config.jwks_url, config.saml_metadata_url, 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", 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", 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 *"#,
            stream.tenant_id, stream.destination_type, stream.destination_url, stream.auth_config, 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 AND tenant_id = $2
               RETURNING *"#,
            stream.id, stream.tenant_id, stream.destination_type, stream.destination_url, stream.auth_config, 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", 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", 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 *"#,
            p.tenant_id, p.entity_type, p.retention_days, p.shred_pii, 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", 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>> {
        Ok(sqlx::query_as!(ApiKey, "SELECT * FROM api_keys WHERE key_hash = $1", 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 *"#,
            k.id, k.tenant_id, k.name, k.key_hash, k.key_prefix, k.scopes, 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", id).execute(&mut *tx).await?;
        tx.commit().await?;
        Ok(())
    }
    async fn update_last_used(&self, id: Uuid) -> Result<()> {
        sqlx::query!("UPDATE api_keys SET last_used_at = NOW() WHERE id = $1", 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", 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 *"#,
            p.id, p.tenant_id, p.field_name, p.priority_config, 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", 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", 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 *"#,
            m.id, m.tenant_id, m.data_source_id, 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", 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 *"#,
            tenant_id, version, plan_hash, raw_yaml, 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", 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)
    }
}
