use crate::models::{ChatMessage, ChatSession, IngestionReceipt};
use anyhow::{Context, Result};
use std::path::Path;
use tokio_postgres::{Client, Config, NoTls};
use tracing::{debug, info, warn};

/// PostgreSQL sink for the ENE chat log schema.
pub struct RdsSink {
    client: Client,
}

impl RdsSink {
    /// Connect to PostgreSQL.  DSN is parsed from standard libpq key=value format.
    ///
    /// RDS requires SSL.  When `sslmode=require` (or `verify-*`) is present in the
    /// DSN the caller should set `PGSSLMODE` or pass a DSN that already includes
    /// `sslmode=require`.  We connect with `NoTls` here so the binary has zero
    /// native-tls/openssl link dependency; in production the caller is expected to
    /// tunnel via an SSL-terminating proxy (RDS IAM auth token + pg_bouncer/SSL
    /// termination), OR the `RDS_DSN` env-var contains `sslmode=disable` for a
    /// dev/tunnel setup.  The feature flag `tls` (future) can replace NoTls.
    pub async fn connect(dsn: &str) -> Result<Self> {
        // Strip sslmode=require so tokio-postgres (NoTls) doesn't reject it.
        // In production, place an SSL-terminating proxy in front or use
        // `tokio-postgres-rustls` crate (requires adding it as a dep).
        let cleaned = dsn
            .split_whitespace()
            .filter(|token| !token.starts_with("sslmode="))
            .collect::<Vec<_>>()
            .join(" ");

        let config: Config = cleaned.parse().context("parse PostgreSQL DSN")?;
        let (client, connection) = config
            .connect(NoTls)
            .await
            .context("connect to RDS (NoTls — see sink.rs comment for TLS upgrade path)")?;
        tokio::spawn(async move {
            if let Err(e) = connection.await {
                warn!("PostgreSQL connection error: {}", e);
            }
        });
        let sink = Self { client };
        sink.init_schema().await?;
        sink.init_tables().await?;
        Ok(sink)
    }

    async fn init_schema(&self) -> Result<()> {
        self.client
            .batch_execute("CREATE SCHEMA IF NOT EXISTS ene")
            .await
            .context("create ene schema")?;
        Ok(())
    }

    /// Create tables and indexes if they do not exist.
    async fn init_tables(&self) -> Result<()> {
        // Canonical DDL lives in sql/ene_chat_schema.sql.
        // Keeping it out of Rust source reduces schema drift across services.
        let ddl_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("sql/ene_chat_schema.sql");
        let ddl = std::fs::read_to_string(&ddl_path)
            .with_context(|| format!("read ENE chat schema DDL at {:?}", ddl_path))?;

        self.client.batch_execute(&ddl).await.context("init DDL")?;
        info!("RDS schema initialized");
        Ok(())
    }

    /// Upsert a session (insert or update on conflict).
    pub async fn upsert_session(&self, s: &ChatSession) -> Result<()> {
        let embedding_str = s.embedding.as_ref().map(|v| {
            format!(
                "[{}]",
                v.iter()
                    .map(|f| f.to_string())
                    .collect::<Vec<_>>()
                    .join(",")
            )
        });
        let meta_json = serde_json::to_value(&s.meta)?;
        self.client
            .execute(
                "INSERT INTO ene.chat_sessions \\
                 (session_id, title, agent, model, \\
                  workspace_fingerprint, workspace_root, fork_parent_session_id, \\
                  compaction_count, compaction_summary, message_count, \\
                  token_input_total, token_output_total, \\
                  created_at_ms, updated_at_ms, first_message_at_ms, last_message_at_ms, \\
                  embedding, meta, receipt, updated_at) \\
                 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17::vector,$18,$19,now()) \\
                 ON CONFLICT (session_id) DO UPDATE SET \\
                  title                  = EXCLUDED.title, \\
                  agent                  = EXCLUDED.agent, \\
                  model                  = EXCLUDED.model, \\
                  workspace_fingerprint  = EXCLUDED.workspace_fingerprint, \\
                  workspace_root         = EXCLUDED.workspace_root, \\
                  fork_parent_session_id = EXCLUDED.fork_parent_session_id, \\
                  compaction_count       = EXCLUDED.compaction_count, \\
                  compaction_summary     = EXCLUDED.compaction_summary, \\
                  message_count          = EXCLUDED.message_count, \\
                  token_input_total      = EXCLUDED.token_input_total, \\
                  token_output_total     = EXCLUDED.token_output_total, \\
                  updated_at_ms          = EXCLUDED.updated_at_ms, \\
                  first_message_at_ms    = EXCLUDED.first_message_at_ms, \\
                  last_message_at_ms     = EXCLUDED.last_message_at_ms, \\
                  embedding              = EXCLUDED.embedding, \\
                  meta                   = EXCLUDED.meta, \\
                  receipt                = EXCLUDED.receipt, \\
                  updated_at             = now()",
                &[
                    &s.session_id,
                    &s.title,
                    &s.agent,
                    &s.model,
                    &s.workspace_fingerprint,
                    &s.workspace_root,
                    &s.fork_parent_session_id,
                    &(s.compaction_count as i32),
                    &s.compaction_summary,
                    &(s.message_count as i32),
                    &s.token_input_total,
                    &s.token_output_total,
                    &s.created_at_ms,
                    &s.updated_at_ms,
                    &s.first_message_at_ms,
                    &s.last_message_at_ms,
                    &embedding_str,
                    &meta_json,
                    &s.receipt,
                ],
            )
            .await
            .context("upsert session")?;
        debug!("upserted session {}", s.session_id);
        Ok(())
    }

    /// Upsert a batch of messages for a session.
    pub async fn upsert_messages(&self, session_id: &str, msgs: &[ChatMessage]) -> Result<()> {
        for msg in msgs {
            let embedding_str = msg.embedding.as_ref().map(|v| {
                format!(
                    "[{}]",
                    v.iter()
                        .map(|f| f.to_string())
                        .collect::<Vec<_>>()
                        .join(",")
                )
            });
            let blocks_json = serde_json::to_value(&msg.blocks)?;
            let tool_calls_json = serde_json::to_value(&msg.tool_calls)?;
            self.client
                .execute(
                    "INSERT INTO ene.chat_messages \\
                     (session_id, message_index, role, blocks, text_content, \\
                      token_input, token_output, token_cache_creation, token_cache_read, \\
                      tool_calls, embedding, receipt_hash, created_at_ms, created_at) \\
                     VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::vector,$12,$13,now()) \\
                     ON CONFLICT (session_id, message_index) DO UPDATE SET \\
                      role                 = EXCLUDED.role, \\
                      blocks               = EXCLUDED.blocks, \\
                      text_content         = EXCLUDED.text_content, \\
                      token_input          = EXCLUDED.token_input, \\
                      token_output         = EXCLUDED.token_output, \\
                      token_cache_creation = EXCLUDED.token_cache_creation, \\
                      token_cache_read     = EXCLUDED.token_cache_read, \\
                      tool_calls           = EXCLUDED.tool_calls, \\
                      embedding            = EXCLUDED.embedding, \\
                      receipt_hash         = EXCLUDED.receipt_hash, \\
                      created_at_ms        = EXCLUDED.created_at_ms",
                    &[
                        &session_id,
                        &msg.message_index,
                        &msg.role,
                        &blocks_json,
                        &msg.text_content,
                        &msg.token_input,
                        &msg.token_output,
                        &msg.token_cache_creation,
                        &msg.token_cache_read,
                        &tool_calls_json,
                        &embedding_str,
                        &msg.receipt_hash,
                        &msg.created_at_ms,
                    ],
                )
                .await
                .with_context(|| {
                    format!(
                        "upsert message idx={} for session {}",
                        msg.message_index, session_id
                    )
                })?;
        }
        info!(
            "upserted {} messages for session {}",
            msgs.len(),
            session_id
        );
        Ok(())
    }

    /// Delete all messages for a session (used before re-ingesting).
    pub async fn delete_messages_for_session(&self, session_id: &str) -> Result<u64> {
        let rows = self
            .client
            .execute(
                "DELETE FROM ene.chat_messages WHERE session_id = $1",
                &[&session_id],
            )
            .await
            .context("delete messages")?;
        Ok(rows)
    }

    /// Write an ingestion receipt.
    pub async fn write_receipt(&self, r: &IngestionReceipt) -> Result<()> {
        let meta_json = serde_json::to_value(&r.meta)?;
        self.client
            .execute(
                "INSERT INTO ene.ingestion_receipts \\
                 (shim_name, status, sha256, record_count, source_path, meta) \\
                 VALUES ($1, $2, $3, $4, $5, $6)",
                &[
                    &r.shim_name,
                    &r.status,
                    &r.sha256,
                    &r.record_count,
                    &r.source_path,
                    &meta_json,
                ],
            )
            .await
            .context("write receipt")?;
        Ok(())
    }

    /// Keyword full-text search across messages.
    pub async fn search_keyword(
        &self,
        query: &str,
        limit: i64,
    ) -> Result<Vec<serde_json::Value>> {
        let rows = self
            .client
            .query(
                "SELECT s.session_id, s.title, s.agent, s.model, \\
                 COUNT(m.id) AS match_count, \\
                 MAX(ts_rank(to_tsvector('english', COALESCE(m.text_content,'')), \\
                             plainto_tsquery('english', $1))) AS rank \\
                 FROM ene.chat_sessions s \\
                 JOIN ene.chat_messages m ON m.session_id = s.session_id \\
                 WHERE to_tsvector('english', COALESCE(m.text_content,'')) \\
                       @@ plainto_tsquery('english', $1) \\
                 GROUP BY s.session_id, s.title, s.agent, s.model \\
                 ORDER BY rank DESC \\
                 LIMIT $2",
                &[&query, &limit],
            )
            .await
            .context("keyword search")?;
        let mut out = Vec::new();
        for row in &rows {
            out.push(serde_json::json!({
                "session_id": row.try_get::<_, String>(0)?,
                "title":      row.try_get::<_, Option<String>>(1)?,
                "agent":      row.try_get::<_, Option<String>>(2)?,
                "model":      row.try_get::<_, Option<String>>(3)?,
                "match_count": row.try_get::<_, i64>(4)?,
                "rank":       row.try_get::<_, f32>(5)?,
            }));
        }
        Ok(out)
    }

    /// Semantic search via embedding cosine similarity (requires pgvector).
    pub async fn search_similar(
        &self,
        embedding: &[f32],
        limit: i64,
    ) -> Result<Vec<serde_json::Value>> {
        let vec_str = format!(
            "[{}]",
            embedding
                .iter()
                .map(|f| f.to_string())
                .collect::<Vec<_>>()
                .join(",")
        );
        let rows = self
            .client
            .query(
                "SELECT session_id, title, agent, model, \\
                 1 - (embedding <=> $1::vector) AS similarity \\
                 FROM ene.chat_sessions \\
                 WHERE embedding IS NOT NULL \\
                 ORDER BY embedding <=> $1::vector \\
                 LIMIT $2",
                &[&vec_str, &limit],
            )
            .await
            .context("similarity search")?;
        let mut out = Vec::new();
        for row in &rows {
            out.push(serde_json::json!({
                "session_id": row.try_get::<_, String>(0)?,
                "title":      row.try_get::<_, Option<String>>(1)?,
                "agent":      row.try_get::<_, Option<String>>(2)?,
                "model":      row.try_get::<_, Option<String>>(3)?,
                "similarity": row.try_get::<_, f32>(4)?,
            }));
        }
        Ok(out)
    }

    /// List recent sessions.
    pub async fn list_sessions(&self, limit: i64) -> Result<Vec<serde_json::Value>> {
        let rows = self
            .client
            .query(
                "SELECT session_id, title, agent, model, message_count, \\
                 token_input_total, token_output_total, created_at_ms, updated_at_ms \\
                 FROM ene.chat_sessions \\
                 ORDER BY updated_at_ms DESC \\
                 LIMIT $1",
                &[&limit],
            )
            .await
            .context("list sessions")?;
        let mut out = Vec::new();
        for row in &rows {
            out.push(serde_json::json!({
                "session_id":        row.try_get::<_, String>(0)?,
                "title":             row.try_get::<_, Option<String>>(1)?,
                "agent":             row.try_get::<_, Option<String>>(2)?,
                "model":             row.try_get::<_, Option<String>>(3)?,
                "message_count":     row.try_get::<_, i32>(4)?,
                "token_input_total": row.try_get::<_, i64>(5)?,
                "token_output_total":row.try_get::<_, i64>(6)?,
                "created_at_ms":     row.try_get::<_, i64>(7)?,
                "updated_at_ms":     row.try_get::<_, i64>(8)?,
            }));
        }
        Ok(out)
    }

    /// Retrieve a single session with all its messages.
    pub async fn get_session(&self, session_id: &str) -> Result<Option<serde_json::Value>> {
        let Some(sess) = self
            .client
            .query_opt(
                "SELECT session_id, title, agent, model, message_count, \\
                 token_input_total, token_output_total, created_at_ms, updated_at_ms, meta \\
                 FROM ene.chat_sessions WHERE session_id = $1",
                &[&session_id],
            )
            .await
            .context("get session")?
        else {
            return Ok(None);
        };

        let msg_rows = self
            .client
            .query(
                "SELECT message_index, role, blocks, text_content, \\
                 token_input, token_output, tool_calls, created_at_ms \\
                 FROM ene.chat_messages \\
                 WHERE session_id = $1 ORDER BY message_index",
                &[&session_id],
            )
            .await
            .context("get messages")?;
        let mut messages = Vec::new();
        for row in &msg_rows {
            messages.push(serde_json::json!({
                "message_index": row.try_get::<_, i32>(0)?,
                "role":          row.try_get::<_, String>(1)?,
                "blocks":        row.try_get::<_, serde_json::Value>(2)?,
                "text_content":  row.try_get::<_, Option<String>>(3)?,
                "token_input":   row.try_get::<_, i64>(4)?,
                "token_output":  row.try_get::<_, i64>(5)?,
                "tool_calls":    row.try_get::<_, serde_json::Value>(6)?,
                "created_at_ms": row.try_get::<_, i64>(7)?,
            }));
        }
        Ok(Some(serde_json::json!({
            "session_id":         sess.try_get::<_, String>(0)?,
            "title":              sess.try_get::<_, Option<String>>(1)?,
            "agent":              sess.try_get::<_, Option<String>>(2)?,
            "model":              sess.try_get::<_, Option<String>>(3)?,
            "message_count":      sess.try_get::<_, i32>(4)?,
            "token_input_total":  sess.try_get::<_, i64>(5)?,
            "token_output_total": sess.try_get::<_, i64>(6)?,
            "created_at_ms":      sess.try_get::<_, i64>(7)?,
            "updated_at_ms":      sess.try_get::<_, i64>(8)?,
            "meta":               sess.try_get::<_, serde_json::Value>(9)?,
            "messages": messages,
        })))
    }
}