// p9.go - 9P filesystem server for toolsrvclient.
package main

import (
	"context"
	"fmt"
	"hash/fnv"
	"io"
	"net"
	"sync"

	"ollie/cmd/toolsrv/internal/server"
	"ollie/virtfs"

	"9fans.net/go/plan9"
)

// fid represents an open file handle.
type fid struct {
	mu       sync.Mutex
	path     string
	qid      plan9.Qid
	mode     uint8
	opened   bool
	entry    virtfs.File
	readBuf  []byte // buffer for read data
	readDone bool   // true if file has been read
	waitBase string // base string for blocking/stream reads
	dirCache []byte
}

// authFid represents an auth file handle for Tauth.
type authFid struct {
	mu            sync.Mutex
	uname         string
	aname         string
	authenticated bool
	token         string
}

// connState holds per-connection state.
type connState struct {
	mu       sync.RWMutex
	fids     map[uint32]*fid
	authFids map[uint32]*authFid
	ctx      context.Context
	uname    string
	srv      *server.Server
	noAuth   bool // skip authentication (for debugging)
}

// serve9P handles a single 9P connection using the virtfs tree.
func serve9P(ctx context.Context, conn net.Conn, tree *virtfs.Tree, srv *server.Server, noAuth bool) {
	connCtx, connCancel := context.WithCancel(ctx)
	defer connCancel()
	defer conn.Close()

	cs := &connState{
		fids:     make(map[uint32]*fid),
		authFids: make(map[uint32]*authFid),
		ctx:      connCtx,
		noAuth:   noAuth,
		srv:      srv,
	}

	responses := make(chan *plan9.Fcall, 16)
	var wg sync.WaitGroup
	var writerWG sync.WaitGroup

	// Writer goroutine
	writerWG.Add(1)
	go func() {
		defer writerWG.Done()
		for resp := range responses {
			if err := plan9.WriteFcall(conn, resp); err != nil {
				logger.Warn("write error: %v", err)
				connCancel()
				conn.Close()
				return
			}
		}
	}()

	stop := func() {
		connCancel()
		conn.Close()
		wg.Wait()
		close(responses)
		writerWG.Wait()
		cs.closeFids()
	}

	for {
		select {
		case <-connCtx.Done():
			stop()
			return
		default:
		}

		fc, err := plan9.ReadFcall(conn)
		if err != nil {
			if err != io.EOF && connCtx.Err() == nil {
				logger.Warn("read error: %v", err)
			}
			stop()
			return
		}

		// Handle requests concurrently to allow Tclunk during blocking operations
		wg.Add(1)
		go func(fc *plan9.Fcall) {
			defer wg.Done()
			resp := handleFcall(cs, fc, tree)
			select {
			case responses <- resp:
			case <-connCtx.Done():
			}
		}(fc)
	}
}

func (cs *connState) closeFids() {
	cs.mu.Lock()
	fids := make([]*fid, 0, len(cs.fids))
	for id, f := range cs.fids {
		delete(cs.fids, id)
		fids = append(fids, f)
	}
	cs.authFids = make(map[uint32]*authFid)
	cs.mu.Unlock()

	for _, f := range fids {
		if f.entry != nil {
			_ = f.entry.Close()
		}
	}
}

func handleFcall(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	switch fc.Type {
	case plan9.Tversion:
		return handleVersion(fc)
	case plan9.Tauth:
		return handleAuth(cs, fc)
	case plan9.Tattach:
		return handleAttach(cs, fc, tree)
	case plan9.Twalk:
		return handleWalk(cs, fc, tree)
	case plan9.Topen:
		return handleOpen(cs, fc, tree)
	case plan9.Tread:
		return handleRead(cs, fc, tree)
	case plan9.Twrite:
		return handleWrite(cs, fc, tree)
	case plan9.Tclunk:
		return handleClunk(cs, fc)
	case plan9.Tstat:
		return handleStat(cs, fc, tree)
	default:
		return errFcall(fc, fmt.Sprintf("unsupported: %d", fc.Type))
	}
}

func errFcall(fc *plan9.Fcall, msg string) *plan9.Fcall {
	return &plan9.Fcall{Type: plan9.Rerror, Tag: fc.Tag, Ename: msg}
}

func handleVersion(fc *plan9.Fcall) *plan9.Fcall {
	msize := fc.Msize
	if msize > 65536 {
		msize = 65536
	}
	return &plan9.Fcall{Type: plan9.Rversion, Tag: fc.Tag, Msize: msize, Version: "9P2000"}
}

func handleAuth(cs *connState, fc *plan9.Fcall) *plan9.Fcall {
	cs.mu.Lock()
	cs.authFids[fc.Afid] = &authFid{
		uname: fc.Uname,
		aname: fc.Aname,
	}
	cs.mu.Unlock()

	return &plan9.Fcall{
		Type: plan9.Rauth,
		Tag:  fc.Tag,
		Aqid: plan9.Qid{Type: plan9.QTAUTH, Vers: 0, Path: uint64(fc.Afid)},
	}
}

func handleAttach(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	// Skip auth check if noAuth is set (for debugging)
	if cs.noAuth {
		cs.mu.Lock()
		cs.uname = fc.Uname
		cs.fids[fc.Fid] = &fid{
			path: "/",
			qid:  plan9.Qid{Type: plan9.QTDIR, Path: 0},
		}
		cs.mu.Unlock()
		return &plan9.Fcall{Type: plan9.Rattach, Tag: fc.Tag, Qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}}
	}

	// Check auth - afid must exist and be authenticated
	cs.mu.RLock()
	af, hasAuth := cs.authFids[fc.Afid]
	cs.mu.RUnlock()

	// NOFID means no auth required (we require auth)
	if fc.Afid == plan9.NOFID {
		return errFcall(fc, "authentication required")
	}

	if !hasAuth {
		return errFcall(fc, "invalid afid")
	}
	af.mu.Lock()
	authenticated := af.authenticated
	af.mu.Unlock()
	if !authenticated {
		return errFcall(fc, "not authenticated")
	}

	cs.mu.Lock()
	cs.uname = fc.Uname
	cs.fids[fc.Fid] = &fid{
		path: "/",
		qid:  plan9.Qid{Type: plan9.QTDIR, Path: 0},
	}
	cs.mu.Unlock()

	return &plan9.Fcall{Type: plan9.Rattach, Tag: fc.Tag, Qid: plan9.Qid{Type: plan9.QTDIR, Path: 0}}
}

// handleAuthWrite handles writing the secret to an auth fid.
func handleAuthWrite(cs *connState, fc *plan9.Fcall, af *authFid) *plan9.Fcall {
	clientSecret := string(fc.Data)
	token, err := cs.srv.Authenticate(clientSecret)
	if err != nil {
		return errFcall(fc, err.Error())
	}

	af.mu.Lock()
	af.authenticated = true
	af.token = token
	af.mu.Unlock()

	return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))}
}

// handleAuthRead handles reading the token from an auth fid.
func handleAuthRead(cs *connState, fc *plan9.Fcall, af *authFid) *plan9.Fcall {
	af.mu.Lock()
	authenticated := af.authenticated
	token := af.token
	af.mu.Unlock()
	if !authenticated {
		return errFcall(fc, "write secret first")
	}

	data := []byte(token + "\n")

	offset := int(fc.Offset)
	if offset >= len(data) {
		return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil}
	}

	end := offset + int(fc.Count)
	if end > len(data) {
		end = len(data)
	}

	return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: data[offset:end]}
}

func handleWalk(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	cs.mu.RLock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.RUnlock()
	if !ok {
		return errFcall(fc, "unknown fid")
	}

	newPath := f.path
	var qids []plan9.Qid

	for _, name := range fc.Wname {
		if name == ".." {
			if newPath != "/" {
				lastSlash := len(newPath) - 1
				for lastSlash > 0 && newPath[lastSlash] != '/' {
					lastSlash--
				}
				if lastSlash == 0 {
					newPath = "/"
				} else {
					newPath = newPath[:lastSlash]
				}
			}
		} else {
			if newPath == "/" {
				newPath = "/" + name
			} else {
				newPath = newPath + "/" + name
			}
		}

		// Check if path exists in tree
		treePath := newPath
		if treePath == "/" {
			treePath = ""
		} else {
			treePath = treePath[1:] // strip leading /
		}

		info, err := tree.Stat(treePath)
		if err != nil {
			if len(qids) == 0 {
				return errFcall(fc, "not found")
			}
			break
		}

		var qt uint8 = plan9.QTFILE
		if info.IsDir() {
			qt = plan9.QTDIR
		}
		qids = append(qids, plan9.Qid{Type: qt, Path: pathHash(newPath)})
	}

	if len(qids) == len(fc.Wname) && fc.Newfid != fc.Fid {
		var lastQid plan9.Qid
		if len(qids) > 0 {
			lastQid = qids[len(qids)-1]
		} else {
			lastQid = f.qid
		}
		cs.mu.Lock()
		cs.fids[fc.Newfid] = &fid{
			path: newPath,
			qid:  lastQid,
		}
		cs.mu.Unlock()
	} else if len(qids) == len(fc.Wname) && len(qids) > 0 {
		cs.mu.Lock()
		f.path = newPath
		f.qid = qids[len(qids)-1]
		cs.mu.Unlock()
	}

	return &plan9.Fcall{Type: plan9.Rwalk, Tag: fc.Tag, Wqid: qids}
}

func handleOpen(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	cs.mu.Lock()
	defer cs.mu.Unlock()

	f, ok := cs.fids[fc.Fid]
	if !ok {
		return errFcall(fc, "unknown fid")
	}

	// Directories don't need an entry
	if f.qid.Type&plan9.QTDIR != 0 {
		f.opened = true
		f.mode = fc.Mode
		return &plan9.Fcall{Type: plan9.Ropen, Tag: fc.Tag, Qid: f.qid, Iounit: 8192}
	}

	path := f.path
	if path == "/" {
		path = ""
	} else {
		path = path[1:] // strip leading /
	}

	entry, err := tree.Open(path)
	if err != nil {
		return errFcall(fc, err.Error())
	}

	f.entry = entry
	f.opened = true
	f.mode = fc.Mode
	f.readBuf = nil
	f.readDone = false

	return &plan9.Fcall{Type: plan9.Ropen, Tag: fc.Tag, Qid: f.qid, Iounit: 8192}
}

func handleRead(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	// Check if reading from auth fid
	cs.mu.RLock()
	af, isAuthFid := cs.authFids[fc.Fid]
	f, isFid := cs.fids[fc.Fid]
	cs.mu.RUnlock()

	if isAuthFid {
		return handleAuthRead(cs, fc, af)
	}

	if !isFid {
		return errFcall(fc, "unknown fid")
	}
	f.mu.Lock()
	opened := f.opened
	qid := f.qid
	entry := f.entry
	f.mu.Unlock()
	if !opened {
		return errFcall(fc, "not open")
	}

	if qid.Type&plan9.QTDIR != 0 {
		return handleDirRead(cs, f, fc, tree)
	}

	if entry == nil {
		return errFcall(fc, "no entry")
	}

	// For rdwr files, check if we've already read
	f.mu.Lock()
	readDone := entry.RdwrMode() && f.readDone
	f.mu.Unlock()
	if readDone {
		return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil}
	}

	// Handle BlockOnce mode: block until content is available, then cache it
	if f.entry.BlockingReadMode() {
		// Serve from cache if we already have content (multi-Tread paging)
		if f.readBuf != nil {
			if int(fc.Offset) >= len(f.readBuf) {
				return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Count: 0}
			}
			return readSlice(fc, f.readBuf)
		}
		// First Tread: block for content
		content, nextBase, err := f.entry.BlockingRead(cs.ctx, f.waitBase)
		if err != nil {
			return errFcall(fc, err.Error())
		}
		if nextBase != "" {
			cs.mu.Lock()
			if ff, ok := cs.fids[fc.Fid]; ok {
				ff.waitBase = nextBase
			}
			cs.mu.Unlock()
		}
		f.readBuf = content
		return readSlice(fc, content)
	}

	// Handle Stream mode: block indefinitely, stream data incrementally
	if f.entry.StreamMode() {
		content, nextBase, err := f.entry.BlockingRead(cs.ctx, f.waitBase)
		if err != nil {
			return errFcall(fc, err.Error())
		}
		if nextBase != "" {
			cs.mu.Lock()
			if ff, ok := cs.fids[fc.Fid]; ok {
				ff.waitBase = nextBase
			}
			cs.mu.Unlock()
		}
		count := int(fc.Count)
		if count > len(content) {
			count = len(content)
		}
		return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Count: uint32(count), Data: content[:count]}
	}

	// Buffer read data if not already done (normal non-blocking read)
	if f.readBuf == nil {
		data, err := f.entry.Read()
		if err != nil && err != io.EOF {
			return errFcall(fc, err.Error())
		}
		f.readBuf = data
	}

	// Handle offset/count
	offset := int(fc.Offset)
	if offset >= len(f.readBuf) {
		f.readDone = true
		return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil}
	}

	end := offset + int(fc.Count)
	if end > len(f.readBuf) {
		end = len(f.readBuf)
	}

	return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: f.readBuf[offset:end]}
}

// readSlice returns a slice of content based on offset/count from fc.
func readSlice(fc *plan9.Fcall, content []byte) *plan9.Fcall {
	var data []byte
	off := int(fc.Offset)
	if off < len(content) {
		end := off + int(fc.Count)
		if end > len(content) {
			end = len(content)
		}
		data = content[off:end]
	}
	return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Count: uint32(len(data)), Data: data}
}

func handleDirRead(cs *connState, f *fid, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	f.mu.Lock()
	path := f.path
	qidPath := f.path
	cached := f.dirCache
	f.mu.Unlock()
	if fc.Offset == 0 || cached == nil {
		if path == "/" {
			path = ""
		} else {
			path = path[1:]
		}

		entries, err := tree.Readdir(path)
		if err != nil {
			return errFcall(fc, err.Error())
		}

		var buf []byte
		for _, e := range entries {
			var qt uint8 = plan9.QTFILE
			mode := plan9.Perm(0644)
			if e.IsDir() {
				qt = plan9.QTDIR
				mode = plan9.Perm(plan9.DMDIR | 0755)
			}
			dir := plan9.Dir{
				Name: e.Name(),
				Qid:  plan9.Qid{Type: qt, Path: pathHash(qidPath + "/" + e.Name())},
				Mode: mode,
				Uid:  "toolsrv",
				Gid:  "toolsrv",
			}
			data, _ := dir.Bytes()
			buf = append(buf, data...)
		}
		f.mu.Lock()
		if fc.Offset == 0 || f.dirCache == nil {
			f.dirCache = buf
		}
		cached = f.dirCache
		f.mu.Unlock()
	}

	offset := int(fc.Offset)
	if offset >= len(cached) {
		return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: nil}
	}

	end := offset + int(fc.Count)
	if end > len(cached) {
		end = len(cached)
	}

	return &plan9.Fcall{Type: plan9.Rread, Tag: fc.Tag, Data: cached[offset:end]}
}

func handleWrite(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	// Check if writing to auth fid
	cs.mu.Lock()
	af, isAuthFid := cs.authFids[fc.Fid]
	f, isFid := cs.fids[fc.Fid]
	cs.mu.Unlock()

	if isAuthFid {
		return handleAuthWrite(cs, fc, af)
	}

	if !isFid {
		return errFcall(fc, "unknown fid")
	}
	f.mu.Lock()
	opened := f.opened
	entry := f.entry
	f.mu.Unlock()
	if !opened {
		return errFcall(fc, "not open")
	}
	if entry == nil {
		return errFcall(fc, "no entry")
	}

	err := entry.Write(fc.Data)
	if err != nil {
		return errFcall(fc, err.Error())
	}

	return &plan9.Fcall{Type: plan9.Rwrite, Tag: fc.Tag, Count: uint32(len(fc.Data))}
}

func handleClunk(cs *connState, fc *plan9.Fcall) *plan9.Fcall {
	cs.mu.Lock()
	f, ok := cs.fids[fc.Fid]
	if ok {
		delete(cs.fids, fc.Fid)
	}
	delete(cs.authFids, fc.Fid)
	cs.mu.Unlock()

	// Close the file handle to cancel any pending Request handlers.
	if ok && f.entry != nil {
		f.entry.Close()
	}

	return &plan9.Fcall{Type: plan9.Rclunk, Tag: fc.Tag}
}

func handleStat(cs *connState, fc *plan9.Fcall, tree *virtfs.Tree) *plan9.Fcall {
	cs.mu.RLock()
	f, ok := cs.fids[fc.Fid]
	cs.mu.RUnlock()
	if !ok {
		return errFcall(fc, "unknown fid")
	}

	path := f.path
	if path == "/" {
		path = ""
	} else {
		path = path[1:]
	}

	info, err := tree.Stat(path)
	if err != nil {
		return errFcall(fc, err.Error())
	}

	var qt uint8 = plan9.QTFILE
	mode := plan9.Perm(0644)
	if info.IsDir() {
		qt = plan9.QTDIR
		mode = plan9.Perm(plan9.DMDIR | 0755)
	}

	name := info.Name()
	if name == "" {
		name = "/"
	}

	dir := plan9.Dir{
		Name: name,
		Qid:  plan9.Qid{Type: qt, Path: pathHash(f.path)},
		Mode: mode,
		Uid:  "toolsrv",
		Gid:  "toolsrv",
	}
	data, _ := dir.Bytes()

	return &plan9.Fcall{Type: plan9.Rstat, Tag: fc.Tag, Stat: data}
}

func pathHash(p string) uint64 {
	h := fnv.New64a()
	h.Write([]byte(p))
	return h.Sum64()
}