package zig

import (
	"fmt"
	"io"
)

type formatter struct {
	w      io.Writer
	line   int // 1-based
	col    int // 1-based, reset to 1 after newline
	indent int // indentation level
}

const indentStr = "    " // 4 spaces per indent

func (f *formatter) Writef(format string, a ...any) {
	s := fmt.Sprintf(format, a...)
	for i, r := range s {
		if r == '\n' {
			f.line++
			f.col = 1
			// If next character is not a newline and not end, and not a closing brace, write indentation
			if i+1 < len(s) && s[i+1] != '\n' && s[i+1] != '}' {
				f.writeIndent()
			}
		} else {
			if f.col == 0 {
				f.col = 1
			} else {
				f.col++
			}
		}
	}
	if _, err := f.w.Write([]byte(s)); err != nil {
		panic(err)
	}
}

func (f *formatter) writeIndent() {
	for i := 0; i < f.indent; i++ {
		if _, err := f.w.Write([]byte(indentStr)); err != nil {
			panic(err)
		}
		f.col += len(indentStr)
	}
}

func Write(w io.Writer, root *Root) (err error) {
	defer func() {
		if r := recover(); r != nil {
			if e, ok := r.(error); ok {
				err = e
			} else {
				panic(r)
			}
		}
	}()
	f := &formatter{w: w, line: 1, col: 1, indent: 0}

	if root.ContainerDocComment != "" {
		f.Writef("//! %s\n\n", root.ContainerDocComment)
	}
	for _, member := range root.ContainerMembers {
		// Only handle Decl for now (fields not needed for hello world)
		if member.Decl != nil {
			writeDecl(f, member.Decl)
		}
	}
	return nil
}

func writeDecl(f *formatter, decl Decl) {
	switch fn := decl.(type) {
	case *FnDecl:
		f.Writef("fn %s(", fn.Name)
		writeParams(f, fn.Params)
		f.Writef(") ")
		writeTypeExpr(f, fn.ReturnType)
		writeBlock(f, fn.Body)
	}
}

func writeParams(f *formatter, params []*ParamDecl) {
	for i, param := range params {
		if i > 0 {
			f.Writef(", ")
		}
		if param.Name != "" {
			f.Writef("%s: ", param.Name)
		}
		writeTypeExpr(f, param.Type)
	}
}

func writeTypeExpr(f *formatter, typ TypeExpr) {
	switch t := typ.(type) {
	case *Identifier:
		f.Writef("%s", t.Name)
	case nil:
		// nothing
	default:
		f.Writef("%v", t)
	}
}

func writeBlock(f *formatter, block *Block) {
	if block == nil {
		f.Writef(";")
		return
	}
	f.Writef(" {\n")
	f.indent++
	for i, stmt := range block.Stmts {
		f.writeIndent()
		writeStmt(f, stmt)
		if i < len(block.Stmts)-1 {
			f.Writef("\n")
		}
	}
	f.indent--
	f.Writef("\n")
	f.writeIndent()
	f.Writef("}\n")
}

func writeStmt(f *formatter, stmt Stmt) {
	switch stmt.(type) {
	case *ReturnStmt:
		f.Writef("return;")
	}
}