import { tool } from '@strands-agents/sdk'
import { z } from 'zod'
import { get, set, del } from 'idb-keyval'
import Graph from 'graphology'
import { bfsFromNode, dfsFromNode } from 'graphology-traversal'
import { dijkstra, edgePathFromNodePath } from 'graphology-shortest-path'
import louvain from 'graphology-communities-louvain'

const KG_KEY = 'careless-v2-kg'
let cached: Graph | null = null

async function load(): Promise<Graph> {
  if (cached) return cached
  const blob = await get<any>(KG_KEY)
  const g = new Graph({ multi: false, allowSelfLoops: true, type: 'directed' })
  if (blob && blob.attributes !== undefined) {
    try { g.import(blob) } catch (e) { console.warn('[kg] import failed, starting fresh:', e) }
  }
  cached = g
  return g
}

async function save(g: Graph): Promise<void> {
  cached = g
  await set(KG_KEY, g.export())
}

// Canonical id: lowercased, whitespace→dash, no punctuation
function nodeId(type: string, name: string): string {
  const norm = (s: string) => s.toLowerCase().trim().replace(/\s+/g, '-').replace(/[^a-z0-9:_-]/g, '')
  return `${norm(type)}:${norm(name)}`
}

function ok(data: Record<string, unknown>) {
  return JSON.stringify({ status: 'success', ...data })
}
function fail(e: unknown) {
  return JSON.stringify({ status: 'error', error: (e as Error).message || String(e) })
}

export const kgAddEntityTool = tool({
  name: 'kg_add_entity',
  description: 'Add or update an entity node. Returns node id + total count.',
  inputSchema: z.object({
    name: z.string(),
    type: z.string().describe('Entity type (person/project/technology/concept/etc.)'),
    properties: z.record(z.string(), z.string()).optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      const id = nodeId(input.type, input.name)
      const created = !g.hasNode(id)
      if (created) {
        g.addNode(id, {
          name: input.name.toLowerCase().trim(),  // normalize on write
          type: input.type.toLowerCase().trim(),
          ...(input.properties || {}),
          created: Date.now(),
        })
      } else {
        g.mergeNodeAttributes(id, { ...(input.properties || {}), updated: Date.now() })
      }
      await save(g)
      return ok({ id, created, total_nodes: g.order, total_edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const kgAddRelationTool = tool({
  name: 'kg_add_relation',
  description: 'Add directed edge between two entities. Auto-creates missing nodes with type="unknown".',
  inputSchema: z.object({
    from: z.string().describe('Source node id (type:name format)'),
    to: z.string().describe('Target node id (type:name format)'),
    relation: z.string().describe('Edge label (e.g. works_on, knows, depends_on)'),
    weight: z.number().optional().describe('Edge weight (default 1)'),
    properties: z.record(z.string(), z.string()).optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      const autoCreated: string[] = []
      for (const raw of [input.from, input.to]) {
        if (!g.hasNode(raw)) {
          const [t, n] = raw.includes(':') ? raw.split(':', 2) : ['unknown', raw]
          g.addNode(raw, { name: n, type: t || 'unknown', created: Date.now(), auto_created: true })
          autoCreated.push(raw)
        }
      }
      const attrs = {
        relation: input.relation,
        weight: input.weight ?? 1,
        ...(input.properties || {}),
        created: Date.now(),
      }
      const key = `${input.from}->${input.to}:${input.relation}`
      if (g.hasEdge(key)) g.mergeEdgeAttributes(key, attrs)
      else g.addEdgeWithKey(key, input.from, input.to, attrs)
      await save(g)
      return ok({ edge: key, auto_created: autoCreated, total_nodes: g.order, total_edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const kgQueryTool = tool({
  name: 'kg_query',
  description: 'Query nodes by case-insensitive name substring and/or type. Optionally include adjacent edges.',
  inputSchema: z.object({
    name: z.string().optional(),
    type: z.string().optional(),
    includeRelations: z.boolean().optional(),
    limit: z.number().optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      const q = input.name?.toLowerCase().trim()
      const typeFilter = input.type?.toLowerCase().trim()
      const matches: any[] = []
      g.forEachNode((id, attrs) => {
        if (typeFilter && String(attrs.type).toLowerCase() !== typeFilter) return
        if (q) {
          const nameStr = String(attrs.name || '').toLowerCase()
          if (!nameStr.includes(q) && !id.toLowerCase().includes(q)) return
        }
        matches.push({ id, ...attrs })
      })
      const limited = matches.slice(0, input.limit ?? 50)
      let edges: any[] | undefined
      if (input.includeRelations) {
        const ids = new Set(limited.map(m => m.id))
        edges = []
        g.forEachEdge((key, attrs, src, tgt) => {
          if (ids.has(src) || ids.has(tgt)) edges!.push({ key, from: src, to: tgt, ...attrs })
        })
      }
      return ok({ count: matches.length, returned: limited.length, nodes: limited, ...(edges ? { edges } : {}) })
    } catch (e) { return fail(e) }
  },
})

export const kgNeighborsTool = tool({
  name: 'kg_neighbors',
  description: '1-hop neighbors of a node. Direction: in/out/both.',
  inputSchema: z.object({
    id: z.string(),
    direction: z.enum(['in', 'out', 'both']).optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      if (!g.hasNode(input.id)) return fail(new Error(`node not found: ${input.id}`))
      const dir = input.direction ?? 'both'
      const out: any[] = []
      const seen = new Set<string>()
      const add = (nid: string, edgeKey: string, role: 'in' | 'out') => {
        const k = nid + '|' + edgeKey
        if (seen.has(k)) return
        seen.add(k)
        out.push({ id: nid, role, edge: edgeKey, ...g.getNodeAttributes(nid), edgeAttrs: g.getEdgeAttributes(edgeKey) })
      }
      if (dir === 'out' || dir === 'both') g.forEachOutEdge(input.id, (k, _a, _s, t) => add(t, k, 'out'))
      if (dir === 'in'  || dir === 'both') g.forEachInEdge(input.id, (k, _a, s) => add(s, k, 'in'))
      return ok({ id: input.id, count: out.length, neighbors: out })
    } catch (e) { return fail(e) }
  },
})

export const kgTraverseTool = tool({
  name: 'kg_traverse',
  description: 'BFS/DFS traversal from a node. Returns visited nodes up to maxDepth (default 3). Optional relation filter.',
  inputSchema: z.object({
    start: z.string().describe('Starting node id'),
    mode: z.enum(['bfs', 'dfs']).optional(),
    maxDepth: z.number().optional(),
    relationFilter: z.string().optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      if (!g.hasNode(input.start)) return fail(new Error(`node not found: ${input.start}`))
      const maxDepth = input.maxDepth ?? 3
      const visited: { id: string; depth: number; attrs: any }[] = []
      const walker = input.mode === 'dfs' ? dfsFromNode : bfsFromNode
      let target: Graph = g
      if (input.relationFilter) {
        target = g.copy()
        const toDrop: string[] = []
        target.forEachEdge((k, a) => { if (a.relation !== input.relationFilter) toDrop.push(k) })
        toDrop.forEach(k => target.dropEdge(k))
      }
      walker(target, input.start, (node, attrs, depth) => {
        if (depth > maxDepth) return true
        visited.push({ id: node, depth, attrs })
        return false
      })
      return ok({ start: input.start, mode: input.mode ?? 'bfs', maxDepth, count: visited.length, visited })
    } catch (e) { return fail(e) }
  },
})

export const kgShortestPathTool = tool({
  name: 'kg_shortest_path',
  description: 'Shortest path between two nodes. Returns hops + total cost. Weighted uses edge.weight attribute via Dijkstra.',
  inputSchema: z.object({
    from: z.string(),
    to: z.string(),
    weighted: z.boolean().optional(),
  }),
  callback: async (input) => {
    try {
      const g = await load()
      if (!g.hasNode(input.from)) return fail(new Error(`node not found: ${input.from}`))
      if (!g.hasNode(input.to)) return fail(new Error(`node not found: ${input.to}`))
      const path: string[] | null = input.weighted
        ? dijkstra.bidirectional(g, input.from, input.to, 'weight')
        : dijkstra.bidirectional(g, input.from, input.to)
      if (!path) return ok({ found: false })
      const edgePath = edgePathFromNodePath(g, path)
      // Compute total cost
      let cost = 0
      for (const ek of edgePath) {
        const w = g.getEdgeAttribute(ek, 'weight')
        cost += typeof w === 'number' ? w : 1
      }
      return ok({
        found: true,
        hops: path.length - 1,
        cost: input.weighted ? cost : path.length - 1,
        weighted: !!input.weighted,
        nodes: path,
        edges: edgePath.map(k => ({ key: k, ...g.getEdgeAttributes(k) })),
      })
    } catch (e) { return fail(e) }
  },
})

export const kgCommunitiesTool = tool({
  name: 'kg_communities',
  description: 'Louvain community detection. Works on an undirected projection of the graph.',
  inputSchema: z.object({}),
  callback: async () => {
    try {
      const g = await load()
      if (g.order === 0) return ok({ count: 0, communities: {} })
      if (g.size === 0) {
        // Each node is its own community
        const groups: Record<string, string[]> = {}
        g.forEachNode(n => { groups[n] = [n] })
        return ok({ count: g.order, communities: groups, note: 'no edges; each node isolated' })
      }
      // Build undirected simple projection — louvain requires undirected
      const ug = new Graph({ type: 'undirected', allowSelfLoops: false, multi: false })
      g.forEachNode((n, a) => ug.addNode(n, { ...a }))
      g.forEachEdge((_k, a, s, t) => {
        if (s === t) return  // skip self-loops
        if (ug.hasEdge(s, t)) {
          // Accumulate weight
          const prev = (ug.getEdgeAttribute(s, t, 'weight') as number) || 1
          ug.setEdgeAttribute(s, t, 'weight', prev + ((a.weight as number) || 1))
        } else {
          ug.addEdge(s, t, { weight: (a.weight as number) || 1 })
        }
      })
      const map = louvain(ug, { getEdgeWeight: 'weight' }) as Record<string, number>
      const groups: Record<string, string[]> = {}
      for (const [node, comm] of Object.entries(map)) {
        const key = String(comm)
        ;(groups[key] ||= []).push(node)
      }
      return ok({ count: Object.keys(groups).length, communities: groups, nodes: g.order, edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const kgStatsTool = tool({
  name: 'kg_stats',
  description: 'Graph stats: node/edge counts, density, type + relation distribution, top-degree nodes.',
  inputSchema: z.object({}),
  callback: async () => {
    try {
      const g = await load()
      const types: Record<string, number> = {}
      const relations: Record<string, number> = {}
      const degrees: { id: string; degree: number; name: string; type: string }[] = []
      g.forEachNode((id, a) => {
        const t = String(a.type || 'unknown')
        types[t] = (types[t] || 0) + 1
        degrees.push({ id, degree: g.degree(id), name: String(a.name || id), type: t })
      })
      g.forEachEdge((_k, a) => {
        const r = String(a.relation || 'unknown')
        relations[r] = (relations[r] || 0) + 1
      })
      degrees.sort((a, b) => b.degree - a.degree)
      const density = g.order > 1 ? g.size / (g.order * (g.order - 1)) : 0
      return ok({
        nodes: g.order, edges: g.size,
        density: +density.toFixed(4),
        types, relations,
        top_nodes: degrees.slice(0, 10),
      })
    } catch (e) { return fail(e) }
  },
})

export const kgDeleteTool = tool({
  name: 'kg_delete',
  description: 'Delete a node and all its edges.',
  inputSchema: z.object({ id: z.string() }),
  callback: async (input) => {
    try {
      const g = await load()
      if (!g.hasNode(input.id)) return fail(new Error(`node not found: ${input.id}`))
      const edges = g.degree(input.id)
      g.dropNode(input.id)
      await save(g)
      return ok({ deleted: input.id, edges_removed: edges, total_nodes: g.order, total_edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const kgClearTool = tool({
  name: 'kg_clear',
  description: 'Clear the entire knowledge graph. Requires confirm=true.',
  inputSchema: z.object({ confirm: z.boolean() }),
  callback: async (input) => {
    try {
      if (!input.confirm) return fail(new Error('pass confirm=true to clear'))
      await del(KG_KEY)
      cached = new Graph({ multi: false, allowSelfLoops: true, type: 'directed' })
      return ok({ cleared: true, total_nodes: 0, total_edges: 0 })
    } catch (e) { return fail(e) }
  },
})

export const kgExportTool = tool({
  name: 'kg_export',
  description: 'Export graph as serializable JSON (graphology format).',
  inputSchema: z.object({}),
  callback: async () => {
    try {
      const g = await load()
      return ok({ graph: g.export(), nodes: g.order, edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const kgImportTool = tool({
  name: 'kg_import',
  description: 'Import graph from graphology JSON (replaces current graph). Pass confirm=true.',
  inputSchema: z.object({
    graph: z.any().describe('Graphology-format JSON (output of kg_export)'),
    confirm: z.boolean(),
  }),
  callback: async (input) => {
    try {
      if (!input.confirm) return fail(new Error('pass confirm=true to replace current graph'))
      const g = new Graph({ multi: false, allowSelfLoops: true, type: 'directed' })
      g.import(input.graph)
      await save(g)
      return ok({ imported: true, total_nodes: g.order, total_edges: g.size })
    } catch (e) { return fail(e) }
  },
})

export const KNOWLEDGE_GRAPH_TOOLS = [
  kgAddEntityTool, kgAddRelationTool,
  kgQueryTool, kgNeighborsTool,
  kgTraverseTool, kgShortestPathTool,
  kgCommunitiesTool, kgStatsTool,
  kgDeleteTool, kgClearTool,
  kgExportTool, kgImportTool,
]

export { load as loadGraph, KG_KEY }