package tools

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	_ "image/png"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/johannSo/codey/internal/llm"
)

// defaultReadLimit is how many lines read_file returns when the model does not
// ask for a window.
const defaultReadLimit = 2000

// ReadRecord is when a file was last read and what its mtime was at the time.
type ReadRecord struct {
	At  time.Time
	Mod time.Time
}

// ReadTracker records that a file was read.
type ReadTracker interface {
	MarkRead(path string, mod time.Time)
}

// FileTracker additionally answers what was read, which is what lets edit_file
// refuse to edit a file the model has not looked at. agent.State implements it.
type FileTracker interface {
	ReadTracker
	LastRead(path string) (ReadRecord, bool)
}

// ReadArgs is the read_file input.
type ReadArgs struct {
	Path   string `json:"path" jsonschema_description:"Path to the file, absolute or relative to the working root"`
	Offset int    `json:"offset,omitempty" jsonschema_description:"1-indexed line to start at; omit to start at the beginning"`
	Limit  int    `json:"limit,omitempty" jsonschema_description:"Maximum number of lines to return; defaults to 2000"`
}

// NewRead builds the read_file tool. root anchors relative paths; track may be
// nil; extra names additional directories that may be read but never written,
// such as a skill's own directory outside the working root.
//
// The tool it returns is a BlockTool: a PNG, JPEG, GIF or WebP comes back as an
// image the model can look at, everything else as numbered text.
func NewRead(root string, track ReadTracker, extra ...string) Tool {
	return readTool{Tool: newReadText(root, track, extra), root: root, track: track, extra: extra}
}

func newReadText(root string, track ReadTracker, extra []string) Tool {
	return New("read_file", strings.TrimSpace(`
Read a file from disk. Returns the contents with 1-indexed line numbers, in the
form "   12\tcode". Use offset and limit to page through a large file rather
than reading it whole. Line numbers are for your reference only — never include
them in text you write back to the file. An image file (png, jpeg, gif, webp)
is returned as a picture you can see — use this to look at a screenshot you
took.`),
		func(_ context.Context, in ReadArgs) (string, error) {
			path, err := resolveWith(root, extra, in.Path)
			if err != nil {
				return "", err
			}
			info, err := os.Stat(path)
			switch {
			case os.IsNotExist(err):
				return "", notFound(in.Path, extra)
			case err != nil:
				return "", err
			case info.IsDir():
				return "", fmt.Errorf("%s is a directory, not a file; use glob or grep to look inside it", in.Path)
			}

			data, err := os.ReadFile(path)
			if err != nil {
				return "", err
			}
			if track != nil {
				track.MarkRead(path, info.ModTime())
			}
			if len(data) == 0 {
				return fmt.Sprintf("%s is empty (0 bytes).", in.Path), nil
			}

			lines := splitLines(string(data))
			start := in.Offset
			if start < 1 {
				start = 1
			}
			if start > len(lines) {
				return "", fmt.Errorf("offset %d is past the end of %s, which has %d lines", in.Offset, in.Path, len(lines))
			}
			limit := in.Limit
			if limit <= 0 {
				limit = defaultReadLimit
			}
			end := start + limit - 1
			if end > len(lines) {
				end = len(lines)
			}

			var b strings.Builder
			for i := start; i <= end; i++ {
				fmt.Fprintf(&b, "%6d\t%s\n", i, capLine(lines[i-1]))
			}
			if end < len(lines) {
				fmt.Fprintf(&b, "\n(showing lines %d-%d of %d; call read_file again with offset=%d for more)", start, end, len(lines), end+1)
			}
			return b.String(), nil
		})
}

// readTool is read_file with the image branch on top of the text tool.
type readTool struct {
	Tool
	root  string
	track ReadTracker
	extra []string
}

// maxImageBytes is the most read_file will hand to the model as an image. The
// API's own limit is a few megabytes; a screenshot is far below that, and a
// file above it is not something to spend the context window on unreduced.
const maxImageBytes = 4 << 20

// ExecuteBlocks routes an image path to the image reader and anything else to
// the ordinary text reader.
func (r readTool) ExecuteBlocks(ctx context.Context, raw json.RawMessage) (Output, error) {
	var in ReadArgs
	if len(raw) > 0 {
		if err := json.Unmarshal(raw, &in); err != nil {
			return Output{}, fmt.Errorf("invalid arguments: %w", err)
		}
	}
	mediaType, isImage := imageMediaType(in.Path)
	if !isImage {
		text, err := r.Tool.Execute(ctx, raw)
		return Output{Text: text}, err
	}

	path, err := resolveWith(r.root, r.extra, in.Path)
	if err != nil {
		return Output{}, err
	}
	info, err := os.Stat(path)
	switch {
	case os.IsNotExist(err):
		return Output{}, notFound(in.Path, r.extra)
	case err != nil:
		return Output{}, err
	case info.IsDir():
		return Output{}, fmt.Errorf("%s is a directory, not a file; use glob or grep to look inside it", in.Path)
	case info.Size() > maxImageBytes:
		return Output{}, fmt.Errorf("%s is %s, too large to send as an image (the cap is %s); downscale or crop it first, for example with ImageMagick or a browser screenshot at a smaller viewport",
			in.Path, humanBytes(info.Size()), humanBytes(maxImageBytes))
	}
	data, err := os.ReadFile(path)
	if err != nil {
		return Output{}, err
	}
	if r.track != nil {
		r.track.MarkRead(path, info.ModTime())
	}

	caption := fmt.Sprintf("%s: %s image, %s", in.Path, imageFormatName(mediaType), humanBytes(int64(len(data))))
	if cfg, _, err := image.DecodeConfig(bytes.NewReader(data)); err == nil {
		caption = fmt.Sprintf("%s: %s image, %d×%d, %s", in.Path, imageFormatName(mediaType), cfg.Width, cfg.Height, humanBytes(int64(len(data))))
	}
	return Output{
		Text: caption,
		Blocks: []llm.ContentBlock{
			llm.TextBlock(caption),
			llm.ImageBlock(mediaType, base64.StdEncoding.EncodeToString(data)),
		},
	}, nil
}

// imageMediaType maps a file name to the media type the API accepts for it.
// SVG is deliberately not here: it is text, and the API would not take it as
// an image anyway.
func imageMediaType(path string) (string, bool) {
	switch strings.ToLower(filepath.Ext(path)) {
	case ".png":
		return "image/png", true
	case ".jpg", ".jpeg":
		return "image/jpeg", true
	case ".gif":
		return "image/gif", true
	case ".webp":
		return "image/webp", true
	}
	return "", false
}

func imageFormatName(mediaType string) string {
	return strings.ToUpper(strings.TrimPrefix(mediaType, "image/"))
}

func humanBytes(n int64) string {
	switch {
	case n >= 1<<20:
		return fmt.Sprintf("%.1f MB", float64(n)/(1<<20))
	case n >= 1<<10:
		return fmt.Sprintf("%d KB", n>>10)
	}
	return fmt.Sprintf("%d bytes", n)
}

// splitLines splits on \n and drops the trailing empty element a final newline
// produces, so a 3-line file reports 3 lines.
func splitLines(s string) []string {
	s = strings.ReplaceAll(s, "\r\n", "\n")
	lines := strings.Split(s, "\n")
	if n := len(lines); n > 0 && lines[n-1] == "" {
		lines = lines[:n-1]
	}
	return lines
}

// notFound reports a missing file, naming the skill directories when there are
// any. A relative path anchors to the working root, so a model following a
// skill's own wording asks for "references/note.md" and misses; this error is
// the only place it can learn where that file actually lives.
func notFound(path string, extra []string) error {
	if len(extra) == 0 || filepath.IsAbs(path) {
		return fmt.Errorf("no such file: %s", path)
	}
	return fmt.Errorf("no such file: %s — relative paths are resolved against the working root. "+
		"If this belongs to a skill, join it onto that skill's directory: %s",
		path, strings.Join(extra, ", "))
}

// resolve anchors a relative path to root and refuses anything outside it.
func resolve(root, path string) (string, error) {
	return resolveWith(root, nil, path)
}

// resolveWith is resolve with additional permitted roots. It is used by
// read_file, whose extra roots let a skill's own directory be read without
// widening what edit_file or write_file may touch.
func resolveWith(root string, extra []string, path string) (string, error) {
	if path == "" {
		return "", fmt.Errorf("path is required")
	}
	return jailWith(root, extra, resolveBase(root, path))
}

// resolveBase is resolve for the search tools, where an omitted path means the
// working root rather than an error. It does not jail; callers that take a
// path from the model must.
func resolveBase(root, path string) string {
	if root == "" {
		root = "."
	}
	if path == "" {
		return root
	}
	if filepath.IsAbs(path) {
		return filepath.Clean(path)
	}
	return filepath.Join(root, path)
}

// jail refuses paths outside the working root. Symlinks are resolved first, on
// the existing part of the path, so a link pointing out of the tree cannot be
// used to step around the check. It is cheap and catches a lot; it is not a
// substitute for a sandbox, since bash can still cd anywhere.
func jail(root, path string) (string, error) {
	return jailWith(root, nil, path)
}

// jailWith is jail with additional permitted roots, used by the read-only
// tools so a skill can read the files it ships with. Each extra root gets the
// same symlink resolution as the working root, so an extra root cannot be used
// to step sideways into the rest of the filesystem.
func jailWith(root string, extra []string, path string) (string, error) {
	if root == "" {
		return path, nil
	}
	// A file being created does not exist yet, so resolve the deepest existing
	// ancestor and re-attach the rest.
	probe, rest := path, ""
	for {
		if resolved, err := filepath.EvalSymlinks(probe); err == nil {
			path = filepath.Join(resolved, rest)
			break
		}
		parent := filepath.Dir(probe)
		if parent == probe {
			path = filepath.Clean(path)
			break
		}
		rest = filepath.Join(filepath.Base(probe), rest)
		probe = parent
	}

	realRoot := realpath(root)
	if contains(realRoot, path) {
		return path, nil
	}
	for _, e := range extra {
		if e != "" && contains(realpath(e), path) {
			return path, nil
		}
	}
	if len(extra) > 0 {
		return "", fmt.Errorf("%s is outside the working root %s and every skill directory; codey only touches files under the root", path, realRoot)
	}
	return "", fmt.Errorf("%s is outside the working root %s; codey only touches files under the root", path, realRoot)
}

// realpath resolves symlinks, falling back to a lexical clean for a path that
// does not exist.
func realpath(p string) string {
	if resolved, err := filepath.EvalSymlinks(p); err == nil {
		return resolved
	}
	return filepath.Clean(p)
}

// contains reports whether path is root or sits beneath it.
func contains(root, path string) bool {
	rel, err := filepath.Rel(root, path)
	if err != nil {
		return false
	}
	return rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))
}