package zig

import (
	"fmt"
	"io"
	"strings"
)

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)
			}
		}
	}()
	sb := &strings.Builder{}
	f := &formatter{w: sb, line: 1, col: 1, indent: 0}

	if root.ContainerDocComment != "" {
		f.writef("//! %s\n\n", root.ContainerDocComment)
	}
	for i, member := range root.ContainerMembers {
		if member.Decl != nil {
			// Only emit a leading newline before a function/global after the first declaration
			if i > 0 {
				f.writef("\n")
			}
			writeDecl(f, member.Decl)
		}
	}
	out := sb.String()
	if len(out) == 0 || out[len(out)-1] != '\n' {
		out += "\n"
	}
	_, err = w.Write([]byte(out))
	return err
}

// writeDecl emits a top-level declaration.
func writeDecl(f *formatter, decl Decl) {
	switch d := decl.(type) {
	case *FnDecl:
		if d.Flags&FnExport != 0 {
			f.writef("pub ")
		}
		f.writef("fn %s(", d.Name)
		writeParams(f, d.Params)
		f.writef(") ")
		writeTypeExpr(f, d.ReturnType)
		writeBlock(f, d.Body)
		f.writef("\n")
	case *GlobalVarDecl:
		if d.Flags&GlobalVarConst != 0 {
			f.writef("const %s = ", d.Name)
		} else {
			f.writef("var %s = ", d.Name)
		}
		writeExpr(f, d.Value)
		f.writef(";\n")
	case *ContainerDecl:
		f.writef("struct ")
		writeStructBody(f, d)
	case *ErrorSetDecl:
		writeExpr(f, d)
	case *TestDecl:
		f.writef("test ")
		if d.Name != "" {
			f.writef(`"%s"`, d.Name)
		}
		writeBlock(f, d.Block)
		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 *PrefixTypeExpr:
		f.writef("%s", t.Op)
		writeTypeExpr(f, t.Base)
	case *ArrayTypeExpr:
		f.writef("[")
		writeExpr(f, t.Size)
		f.writef("]")
		writeTypeExpr(f, t.Elem)
	case *ErrorUnionTypeExpr:
		writeTypeExpr(f, t.ErrSet)
		f.writef("!")
		writeTypeExpr(f, t.Type)
	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++
	for _, stmt := range block.Stmts {
		f.writeIndent()
		writeStmt(f, stmt)
		f.writef("\n")
	}
	f.indent--
	f.writeIndent()
	f.writef("}")
}

// 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(";")
	case *VarDeclStmt:
		if s.Const {
			f.writef("const ")
		} else {
			f.writef("var ")
		}
		for i, name := range s.Pattern.Names {
			if i > 0 {
				f.writef(", ")
			}
			f.writef("%s", name)
		}
		if s.Type != nil {
			f.writef(": ")
			writeTypeExpr(f, s.Type)
		}
		if s.Value != nil {
			f.writef(" = ")
			writeExpr(f, s.Value)
		}
		f.writef(";")
	case *BlockStmt:
		writeBlock(f, s.Block)
	case *IfStmt:
		f.writef("if (")
		writeExpr(f, s.Cond)
		f.writef(")")
		
		// Handle payload if present
		if s.Payload != nil {
			f.writef(" ")
			writePayload(f, s.Payload)
		}

		// Always write the then branch as a block
		if block, ok := s.Then.(*BlockStmt); ok {
			writeBlock(f, block.Block)
		} else {
			f.writef(" ")
			writeStmt(f, s.Then)
		}
		if s.Else != nil {
			f.writef(" else")
			if block, ok := s.Else.(*BlockStmt); ok {
				writeBlock(f, block.Block)
			} else {
				f.writef(" ")
				writeStmt(f, s.Else)
			}
		}
	case *LoopStmt:
		if s.Kind == "while" {
			f.writef("while (")
			if wp, ok := s.Prefix.(*WhilePrefix); ok {
				writeExpr(f, wp.Cond)
				if wp.Continue != nil {
					f.writef(") : (")
					writeExpr(f, wp.Continue)
				}
			}
			f.writef(")")
			// Always write the body as a block
			if block, ok := s.Body.(*BlockStmt); ok {
				writeBlock(f, block.Block)
			} else {
				f.writef(" ")
				writeStmt(f, s.Body)
			}
		} else if s.Kind == "for" {
			f.writef("for (")
			if fp, ok := s.Prefix.(*ForPrefix); ok {
				for i, arg := range fp.Args {
					if i > 0 {
						f.writef(", ")
					}
					writeExpr(f, arg.Expr)
					if arg.From != nil {
						if lit, ok := arg.From.(*Literal); ok && lit.Value == "" {
							f.writef("..")
						} else {
							f.writef("...")
							writeExpr(f, arg.From)
						}
					}
				}
				f.writef(")")
				if fp.Payload != nil {
					f.writef(" ")
					writePayload(f, fp.Payload)
				}
			}
			// Always write the body as a block
			if block, ok := s.Body.(*BlockStmt); ok {
				writeBlock(f, block.Block)
			} else {
				f.writef(" ")
				writeStmt(f, s.Body)
			}
		}
		case *DeferStmt:
			if s.ErrDefer {
				f.writef("errdefer")
			} else {
				f.writef("defer")
			}
			f.writef(" ")
			writeStmt(f, s.Stmt)
		case *BreakStmt:
			f.writef("break")
			if s.Label != "" {
				f.writef(" :%s", s.Label)
			}
			if s.Value != nil {
				f.writef(" ")
				writeExpr(f, s.Value)
			}
			f.writef(";")
		case *ContinueStmt:
			f.writef("continue")
			if s.Label != "" {
				f.writef(" :%s", s.Label)
			}
			f.writef(";")
		case *SwitchStmt:
			f.writef("switch (")
			writeExpr(f, s.Cond)
			f.writef(") {\n")
			f.indent++
			for _, prong := range s.Prongs {
				f.writeIndent()
				writeSwitchProng(f, prong)
				f.writef("\n")
			}
			f.indent--
			f.writeIndent()
			f.writef("}")
		}
	}

// writePayload emits a payload (|x|, |*x|, |*x, y|, etc).
func writePayload(f *formatter, payload *Payload) {
	f.writef("|")
	for i, name := range payload.Names {
		if i > 0 {
			f.writef(", ")
		}
		if payload.Pointers[i] {
			f.writef("*")
		}
		f.writef("%s", name)
	}
	f.writef("|")
}

// writeSwitchProng emits a switch prong.
func writeSwitchProng(f *formatter, prong *SwitchProng) {
	for i, c := range prong.Cases {
		if i > 0 {
			f.writef(", ")
		}
		if c.IsElse {
			f.writef("else")
		} else {
			writeExpr(f, c.Expr)
			if c.To != nil {
				f.writef("...")
				writeExpr(f, c.To)
			}
		}
	}
	f.writef(" => ")
	// Check if the expression is actually a statement (like return or break)
	if stmt, ok := prong.Expr.(Stmt); ok {
		// If it's a block, write it directly without the leading space
		if blockStmt, isBlock := stmt.(*BlockStmt); isBlock {
			f.writef("{\n")
			f.indent++
			for _, s := range blockStmt.Block.Stmts {
				f.writeIndent()
				writeStmt(f, s)
				f.writef("\n")
			}
			f.indent--
			f.writeIndent()
			f.writef("},")
		} else {
			// For single statements, write without the semicolon
			switch s := stmt.(type) {
			case *ReturnStmt:
				f.writef("return")
				if s.Value != nil {
					f.writef(" ")
					writeExpr(f, s.Value)
				}
			case *BreakStmt:
				f.writef("break")
				if s.Label != "" {
					f.writef(" :%s", s.Label)
				}
				if s.Value != nil {
					f.writef(" ")
					writeExpr(f, s.Value)
				}
			case *ContinueStmt:
				f.writef("continue")
				if s.Label != "" {
					f.writef(" :%s", s.Label)
				}
			case *ExprStmt:
				writeExpr(f, s.Expr)
			default:
				writeStmt(f, stmt)
			}
			f.writef(",")
		}
	} else {
		writeExpr(f, prong.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(`"%v"`, e.Value)
		default:
			f.writef("%v", e.Value)
		}
	case *InitListExpr:
		if e.Empty {
			f.writef(".{}")
		} else if len(e.Values) > 0 {
			if len(e.Values) == 1 {
				f.writef(".{")
				writeExpr(f, e.Values[0])
				f.writef("}")
			} else {
				f.writef(".{ ")
				for i, v := range e.Values {
					if i > 0 {
						f.writef(", ")
					}
					writeExpr(f, v)
				}
				f.writef(" }")
			}
		}
	case *ContainerDecl:
		if e.Kind == "struct" {
			f.writef("struct ")
			writeStructBody(f, e)
		} else {
			panic("not implemented: " + e.Kind)
		}
	case *TryExpr:
		f.writef("try ")
		writeExpr(f, e.Expr)
	case *BinaryExpr:
		writeExpr(f, e.Left)
		f.writef(" %s ", e.Op)
		writeExpr(f, e.Right)
	case *UnaryExpr:
		f.writef("%s", e.Op)
		writeExpr(f, e.Expr)
	case *IndexExpr:
		writeExpr(f, e.Receiver)
		f.writef("[")
		writeExpr(f, e.Index)
		f.writef("]")
	case *UnreachableExpr:
		f.writef("unreachable")
	case *ErrorSetDecl:
		f.writef("error{\n")
		f.indent++
		for _, name := range e.Names {
			f.writeIndent()
			f.writef("%s,\n", name)
		}
		f.indent--
		f.writeIndent()
		f.writef("}")
	}
}

// writeStructBody emits the body of a struct/union/enum/opaque declaration.
func writeStructBody(f *formatter, decl *ContainerDecl) {
	f.writef("{\n")
	f.indent++
	for _, member := range decl.Fields {
		if member.Field != nil {
			// Field or const
			if member.Field.Type == nil && member.Field.Value != nil {
				// const field
				f.writeIndent()
				f.writef("const %s = ", member.Field.Name)
				writeExpr(f, member.Field.Value)
				f.writef(";\n")
			} else {
				// regular field
				f.writeIndent()
				f.writef("%s: ", member.Field.Name)
				writeTypeExpr(f, member.Field.Type)
				if member.Field.Value != nil {
					f.writef(" = ")
					writeExpr(f, member.Field.Value)
				}
				f.writef(",\n")
			}
		} else if member.Decl != nil {
			// Method or nested decl
			f.writef("\n")
			f.writeIndent()
			writeDecl(f, member.Decl)
		}
	}
	f.indent--
	f.writeIndent()
	f.writef("}")
}