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
}

// indentStr defines the string used for each indentation level (4 spaces).
const indentStr = "    "

// Writef writes formatted text to the underlying writer and updates line/col counters.
// It also handles indentation after newlines when appropriate.
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
			// After a newline, write indentation for the next line unless it's a closing brace or another newline.
			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)
	}
}

// writeIndent writes the current indentation level to the output.
// Call this at the start of a new line before writing statements or closing braces.
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)
	}
}

// Write is the entry point for formatting a Zig AST.
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
}

// writeDecl emits a top-level declaration.
func writeDecl(f *formatter, decl Decl) {
	switch d := decl.(type) {
	case *FnDecl:
		f.Writef("\nfn %s(", d.Name)
		writeParams(f, d.Params)
		f.Writef(") ")
		writeTypeExpr(f, d.ReturnType)
		writeBlock(f, d.Body)
	case *GlobalVarDecl:
		if d.Const {
			f.Writef("const %s = ", d.Name)
		} else {
			f.Writef("var %s = ", d.Name)
		}
		writeExpr(f, d.Value)
		f.Writef(";\n")
	}
}

// writeParams emits function parameters, separated by commas.
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)
	}
}

// writeTypeExpr emits a type expression.
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)
	}
}

// writeBlock emits a block, handling indentation for statements and the closing brace.
func writeBlock(f *formatter, block *Block) {
	if block == nil {
		f.Writef(";")
		return
	}
	f.Writef(" {\n")
	f.indent++ // Increase indentation for block contents.
	for i, stmt := range block.Stmts {
		f.writeIndent() // Indent each statement.
		writeStmt(f, stmt)
		if i < len(block.Stmts)-1 {
			f.Writef("\n")
		}
	}
	f.indent-- // Decrease indentation before closing brace.
	f.Writef("\n")
	f.writeIndent() // Indent the closing brace.
	f.Writef("}\n")
}

// writeStmt emits a statement. Indentation is handled by the caller (writeBlock).
func writeStmt(f *formatter, stmt Stmt) {
	switch s := stmt.(type) {
	case *ReturnStmt:
		f.Writef("return")
		if s.Value != nil {
			f.Writef(" ")
			writeExpr(f, s.Value)
		}
		f.Writef(";")
	case *ExprStmt:
		writeExpr(f, s.Expr)
		f.Writef(";")
	}
}

// writeExpr emits an expression.
func writeExpr(f *formatter, expr Expr) {
	switch e := expr.(type) {
	case *Identifier:
		f.Writef("%s", e.Name)
	case *CallExpr:
		writeExpr(f, e.Fun)
		f.Writef("(")
		for i, arg := range e.Args {
			if i > 0 {
				f.Writef(", ")
			}
			writeExpr(f, arg)
		}
		f.Writef(")")
	case *FieldAccessExpr:
		writeExpr(f, e.Receiver)
		f.Writef(".%s", e.Field)
	case *Literal:
		switch e.Kind {
		case "string":
			f.Writef("%q", e.Value)
		default:
			f.Writef("%v", e.Value)
		}
	case *InitListExpr:
		if e.Empty {
			f.Writef(".{}")
		} else {
			f.Writef(".{")
			// TODO
			f.Writef("}")
		}
	}
}