use anyhow::{Context, Result};
use rusqlite::{params, Connection, OptionalExtension};
use serde::Serialize;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MemoryEdgeType {
    Supersedes,
    Duplicates,
    Conflicts,
    DerivedFrom,
    MergedInto,
    SplitFrom,
}

impl MemoryEdgeType {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Supersedes => "supersedes",
            Self::Duplicates => "duplicates",
            Self::Conflicts => "conflicts",
            Self::DerivedFrom => "derived_from",
            Self::MergedInto => "merged_into",
            Self::SplitFrom => "split_from",
        }
    }
}

#[derive(Debug, Clone, Copy, Default)]
pub struct MemoryEdgeWriteContext<'a> {
    pub state_key_id: Option<i64>,
    pub source_candidate_id: Option<i64>,
    pub evidence_event_ids: &'a [i64],
    pub source_operation_id: Option<i64>,
    pub confidence: Option<f64>,
    pub reason: Option<&'a str>,
}

#[derive(Debug, Clone, PartialEq)]
pub struct MemoryEdgeInput<'a> {
    pub edge_type: MemoryEdgeType,
    pub from_memory_id: Option<i64>,
    pub to_memory_id: Option<i64>,
    pub state_key_id: Option<i64>,
    pub source_candidate_id: Option<i64>,
    pub evidence_event_ids: &'a [i64],
    pub source_operation_id: Option<i64>,
    pub confidence: Option<f64>,
    pub reason: Option<&'a str>,
}

pub fn insert_memory_edge(conn: &Connection, input: &MemoryEdgeInput<'_>) -> Result<i64> {
    let now = chrono::Utc::now().timestamp();
    let evidence_event_ids = if input.evidence_event_ids.is_empty() {
        None
    } else {
        Some(
            serde_json::to_string(input.evidence_event_ids)
                .context("serialize memory edge evidence event ids")?,
        )
    };
    conn.execute(
        "INSERT INTO memory_edges
         (edge_type, from_memory_id, to_memory_id, state_key_id, source_candidate_id,
          evidence_event_ids, source_operation_id, confidence, reason, created_at_epoch)
         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
        params![
            input.edge_type.as_str(),
            input.from_memory_id,
            input.to_memory_id,
            input.state_key_id,
            input.source_candidate_id,
            evidence_event_ids.as_deref(),
            input.source_operation_id,
            input.confidence,
            input.reason,
            now
        ],
    )
    .context("insert memory edge")?;
    Ok(conn.last_insert_rowid())
}

pub fn insert_replacement_edges(
    conn: &Connection,
    edge_type: MemoryEdgeType,
    from_memory_ids: &[i64],
    to_memory_id: i64,
    context: MemoryEdgeWriteContext<'_>,
) -> Result<usize> {
    let state_key_id = context
        .state_key_id
        .or(memory_state_key_id(conn, to_memory_id)?);
    let mut seen = std::collections::HashSet::with_capacity(from_memory_ids.len());
    let mut inserted = 0usize;
    for from_memory_id in from_memory_ids
        .iter()
        .copied()
        .filter(|id| *id != to_memory_id && seen.insert(*id))
    {
        insert_memory_edge(
            conn,
            &MemoryEdgeInput {
                edge_type,
                from_memory_id: Some(from_memory_id),
                to_memory_id: Some(to_memory_id),
                state_key_id,
                source_candidate_id: context.source_candidate_id,
                evidence_event_ids: context.evidence_event_ids,
                source_operation_id: context.source_operation_id,
                confidence: context.confidence,
                reason: context.reason,
            },
        )?;
        inserted += 1;
    }
    Ok(inserted)
}

pub fn insert_supersedes_edges(
    conn: &Connection,
    from_memory_ids: &[i64],
    to_memory_id: i64,
    context: MemoryEdgeWriteContext<'_>,
) -> Result<usize> {
    insert_replacement_edges(
        conn,
        MemoryEdgeType::Supersedes,
        from_memory_ids,
        to_memory_id,
        context,
    )
}

pub fn insert_merged_into_edges(
    conn: &Connection,
    from_memory_ids: &[i64],
    to_memory_id: i64,
    context: MemoryEdgeWriteContext<'_>,
) -> Result<usize> {
    insert_replacement_edges(
        conn,
        MemoryEdgeType::MergedInto,
        from_memory_ids,
        to_memory_id,
        context,
    )
}

fn memory_state_key_id(conn: &Connection, memory_id: i64) -> Result<Option<i64>> {
    Ok(conn
        .query_row(
            "SELECT state_key_id FROM memories WHERE id = ?1",
            [memory_id],
            |row| row.get::<_, Option<i64>>(0),
        )
        .optional()
        .with_context(|| format!("load state_key_id for memory edge target id={memory_id}"))?
        .flatten())
}

#[derive(Debug, Clone, Serialize, PartialEq)]
pub struct MemoryEdgeSummary {
    pub incoming_count: usize,
    pub outgoing_count: usize,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub incoming: Vec<MemoryEdgeReference>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub outgoing: Vec<MemoryEdgeReference>,
}

impl MemoryEdgeSummary {
    pub fn has_edges(&self) -> bool {
        self.incoming_count > 0 || self.outgoing_count > 0
    }
}

#[derive(Debug, Clone, Serialize, PartialEq)]
pub struct MemoryEdgeReference {
    pub id: i64,
    pub edge_type: String,
    pub from_memory_id: Option<i64>,
    pub to_memory_id: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub state_key_id: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source_candidate_id: Option<i64>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    pub evidence_event_ids: Vec<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub source_operation_id: Option<i64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub confidence: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reason: Option<String>,
    pub created_at_epoch: i64,
}

pub fn load_memory_edge_summary(conn: &Connection, memory_id: i64) -> Result<MemoryEdgeSummary> {
    let incoming_count = count_edges(conn, "to_memory_id", memory_id)?;
    let outgoing_count = count_edges(conn, "from_memory_id", memory_id)?;
    Ok(MemoryEdgeSummary {
        incoming_count,
        outgoing_count,
        incoming: load_edge_refs(conn, "to_memory_id", memory_id)?,
        outgoing: load_edge_refs(conn, "from_memory_id", memory_id)?,
    })
}

fn count_edges(conn: &Connection, column: &str, memory_id: i64) -> Result<usize> {
    let sql = format!("SELECT COUNT(*) FROM memory_edges WHERE {column} = ?1");
    let count: i64 = conn.query_row(&sql, [memory_id], |row| row.get(0))?;
    Ok(count as usize)
}

fn load_edge_refs(
    conn: &Connection,
    column: &str,
    memory_id: i64,
) -> Result<Vec<MemoryEdgeReference>> {
    let sql = format!(
        "SELECT id, edge_type, from_memory_id, to_memory_id, state_key_id,
                source_candidate_id, evidence_event_ids, source_operation_id,
                confidence, reason, created_at_epoch
         FROM memory_edges
         WHERE {column} = ?1
         ORDER BY created_at_epoch DESC, id DESC
         LIMIT 25"
    );
    let mut stmt = conn.prepare(&sql)?;
    let rows = stmt.query_map([memory_id], |row| {
        let evidence_json: Option<String> = row.get(6)?;
        let evidence_event_ids = match evidence_json {
            Some(json) => serde_json::from_str::<Vec<i64>>(&json).map_err(|err| {
                rusqlite::Error::FromSqlConversionFailure(
                    6,
                    rusqlite::types::Type::Text,
                    Box::new(err),
                )
            })?,
            None => Vec::new(),
        };
        Ok(MemoryEdgeReference {
            id: row.get(0)?,
            edge_type: row.get(1)?,
            from_memory_id: row.get(2)?,
            to_memory_id: row.get(3)?,
            state_key_id: row.get(4)?,
            source_candidate_id: row.get(5)?,
            evidence_event_ids,
            source_operation_id: row.get(7)?,
            confidence: row.get(8)?,
            reason: row.get(9)?,
            created_at_epoch: row.get(10)?,
        })
    })?;
    crate::db::query::collect_rows(rows).context("load memory edge references")
}
