import { z } from "zod"
import { parse, type TemplateVariable } from "./parser"
import { renderDir } from "./render"

const zodTypeMap: Record<string, () => z.ZodType> = {
  string: () => z.string(),
  number: () => z.number(),
  boolean: () => z.boolean(),
}

type ShapeNode = {
  type?: string
  children: Map<string, ShapeNode>
}

function buildTree(variables: TemplateVariable[]): ShapeNode {
  const root: ShapeNode = { children: new Map() }
  for (const v of variables) {
    const segments = v.path.split(".")
    let node = root
    for (let i = 0; i < segments.length; i++) {
      const seg = segments[i]!
      if (!node.children.has(seg)) {
        node.children.set(seg, { children: new Map() })
      }
      node = node.children.get(seg)!
      if (i === segments.length - 1) {
        node.type = v.type
      }
    }
  }
  return root
}

function nodeToSchema(node: ShapeNode): z.ZodType {
  if (node.children.size === 0 && node.type) {
    const factory = zodTypeMap[node.type]
    if (!factory) throw new Error(`Unsupported type: ${node.type}`)
    return factory()
  }
  const shape: Record<string, z.ZodType> = {}
  for (const [key, child] of node.children) {
    shape[key] = nodeToSchema(child)
  }
  return z.object(shape)
}

function inferSchema(variables: TemplateVariable[]): z.ZodType {
  return nodeToSchema(buildTree(variables))
}

interface Stringable {
  toString: () => string
}

interface Issue {
  message: string,
  path: Stringable
}

export class SchemaMismatchError extends Error {
  constructor(issue: Issue) {
    super(`Shema doesnt match used template variables: ${issue.path}: ${issue.message}` )
  }
}

function validateSchemaMatchesTemplates(userSchema: z.ZodType, variables: TemplateVariable[]) {
  if (userSchema._zod.def.type !== "object") {
    throw new SchemaMismatchError({ path: "", message: "Schema must be a z.object()" })
  }

  // Collect all expected paths: leaf variables + intermediate segments must be objects
  const expected = new Map<string, string>()
  for (const v of variables) {
    const segments = v.path.split(".")
    for (let i = 0; i < segments.length - 1; i++) {
      const intermediate = segments.slice(0, i + 1).join(".")
      if (!expected.has(intermediate)) {
        expected.set(intermediate, "object")
      }
    }
    expected.set(v.path, v.type)
  }

  for (const [path, expectedType] of expected) {
    const segments = path.split(".")
    let current = userSchema
    let currentPath = ""
    for (const seg of segments) {
      currentPath = currentPath ? `${currentPath}.${seg}` : seg
      if (current._zod.def.type !== "object") {
        throw new SchemaMismatchError({ path: currentPath, message: `expected z.object() but schema has z.${current._zod.def.type as string}()` })
      }
      const shape = (current as z.ZodObject<any>).shape
      if (!(seg in shape)) {
        throw new SchemaMismatchError({ path: currentPath, message: `missing in schema` })
      }
      current = shape[seg]
    }
    const actual = current._zod.def.type as string
    if (actual !== expectedType) {
      throw new SchemaMismatchError({ path, message: `expected z.${expectedType}() but schema has z.${actual}()` })
    }
  }
}

export const initRenderer = (dirPath: string) => {
  const variables = parse(dirPath)

  const createRenderer = <S extends z.ZodType>(userSchema: S) => {
    validateSchemaMatchesTemplates(userSchema, variables)
    return (targetPath: string, context: z.infer<S>) => {
      userSchema.parse(context)
      renderDir(dirPath, targetPath, context as Record<string, unknown>)
    }
  }

  return createRenderer 
}