package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"html/template"
	"io"
	"io/fs"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// Message represents one message in the conversation.
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// ChatCompletionRequest is the payload sent to the OpenAI API.
type ChatCompletionRequest struct {
	Model    string    `json:"model"`
	Messages []Message `json:"messages"`
}

// ChatCompletionResponseChoice is a single reply (choice) from the API.
type ChatCompletionResponseChoice struct {
	Message Message `json:"message"`
}

// ChatCompletionResponse represents the API response.
type ChatCompletionResponse struct {
	Choices []ChatCompletionResponseChoice `json:"choices"`
	Error   *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error,omitempty"`
}

var (
	openaiAPIKey string
	indexTmpl    *template.Template
	workDir      string // Path to the workspace directory.
	absWorkDir   string // Absolute workspace path.
)

func init() {
	openaiAPIKey = os.Getenv("OPENAI_API_KEY")
	if openaiAPIKey == "" {
		log.Fatal("Environment variable OPENAI_API_KEY is not set")
	}
	workDir = os.Getenv("WORK_DIR")
	if workDir == "" {
		workDir = "./workspace"
	}
	var err error
	absWorkDir, err = filepath.Abs(workDir)
	if err != nil {
		log.Fatalf("Error getting absolute path for workspace: %v", err)
	}
	if err := os.MkdirAll(absWorkDir, os.ModePerm); err != nil {
		log.Fatalf("Error creating workspace directory: %v", err)
	}

	// Load the HTML template from an external file.
	tmplBytes, err := os.ReadFile("index.html")
	if err != nil {
		log.Fatalf("Error reading index.html: %v", err)
	}
	indexTmpl, err = template.New("index").Parse(string(tmplBytes))
	if err != nil {
		log.Fatalf("Error parsing index.html template: %v", err)
	}
}

func main() {
	http.HandleFunc("/", indexHandler)
	http.HandleFunc("/chat", chatHandler)
	http.HandleFunc("/title", titleHandler)

	// File system endpoints (sandboxed within absWorkDir)
	http.HandleFunc("/files", filesListHandler) // GET: list files.
	http.HandleFunc("/file", fileHandler)       // GET for read; POST for write.

	addr := ":8080"
	log.Println("Server starting on", addr)
	log.Println("Workspace directory:", absWorkDir)
	log.Fatal(http.ListenAndServe(addr, nil))
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	if err := indexTmpl.Execute(w, nil); err != nil {
		http.Error(w, "Unable to load page", http.StatusInternalServerError)
	}
}

func chatHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST allowed", http.StatusMethodNotAllowed)
		return
	}
	var clientRequest struct {
		Messages []Message `json:"messages"`
	}
	if err := json.NewDecoder(r.Body).Decode(&clientRequest); err != nil {
		http.Error(w, "Invalid JSON: "+err.Error(), http.StatusBadRequest)
		return
	}
	apiRequestData := ChatCompletionRequest{
		Model:    "o3-mini",
		Messages: clientRequest.Messages,
	}
	requestBytes, err := json.Marshal(apiRequestData)
	if err != nil {
		http.Error(w, "Error marshalling request: "+err.Error(), http.StatusInternalServerError)
		return
	}
	client := http.Client{Timeout: 60 * time.Second}
	req, err := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", bytes.NewBuffer(requestBytes))
	if err != nil {
		http.Error(w, "Failed to create request: "+err.Error(), http.StatusInternalServerError)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+openaiAPIKey)
	apiResponse, err := client.Do(req)
	if err != nil {
		http.Error(w, "Error calling OpenAI API: "+err.Error(), http.StatusInternalServerError)
		return
	}
	defer apiResponse.Body.Close()
	bodyBytes, err := io.ReadAll(apiResponse.Body)
	if err != nil {
		http.Error(w, "Error reading API response: "+err.Error(), http.StatusInternalServerError)
		return
	}
	var completionResponse ChatCompletionResponse
	if err := json.Unmarshal(bodyBytes, &completionResponse); err != nil {
		http.Error(w, "Error parsing API response: "+err.Error(), http.StatusInternalServerError)
		return
	}
	if completionResponse.Error != nil {
		http.Error(w, "OpenAI API error: "+completionResponse.Error.Message, http.StatusInternalServerError)
		return
	}
	if len(completionResponse.Choices) == 0 {
		http.Error(w, "No choices returned from OpenAI API", http.StatusInternalServerError)
		return
	}
	responsePayload := map[string]string{
		"reply": completionResponse.Choices[0].Message.Content,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(responsePayload)
}

func titleHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST allowed", http.StatusMethodNotAllowed)
		return
	}
	var clientRequest struct {
		Messages []Message `json:"messages"`
	}
	if err := json.NewDecoder(r.Body).Decode(&clientRequest); err != nil {
		http.Error(w, "Invalid JSON: "+err.Error(), http.StatusBadRequest)
		return
	}
	titleRequestMessages := []Message{
		{
			Role:    "system",
			Content: "You are an AI that generates conversation titles. Based on the conversation, please provide a short, descriptive title (five words maximum) with no extra commentary.",
		},
	}
	titleRequestMessages = append(titleRequestMessages, clientRequest.Messages...)
	apiRequestData := ChatCompletionRequest{
		Model:    "o3-mini",
		Messages: titleRequestMessages,
	}
	requestBytes, err := json.Marshal(apiRequestData)
	if err != nil {
		http.Error(w, "Error marshalling title request: "+err.Error(), http.StatusInternalServerError)
		return
	}
	client := http.Client{Timeout: 30 * time.Second}
	req, err := http.NewRequest("POST", "https://api.openai.com/v1/chat/completions", bytes.NewBuffer(requestBytes))
	if err != nil {
		http.Error(w, "Failed to create title request: "+err.Error(), http.StatusInternalServerError)
		return
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+openaiAPIKey)
	apiResponse, err := client.Do(req)
	if err != nil {
		http.Error(w, "Error calling OpenAI API for title: "+err.Error(), http.StatusInternalServerError)
		return
	}
	defer apiResponse.Body.Close()
	bodyBytes, err := io.ReadAll(apiResponse.Body)
	if err != nil {
		http.Error(w, "Error reading title API response: "+err.Error(), http.StatusInternalServerError)
		return
	}
	var completionResponse ChatCompletionResponse
	if err := json.Unmarshal(bodyBytes, &completionResponse); err != nil {
		http.Error(w, "Error parsing title API response: "+err.Error(), http.StatusInternalServerError)
		return
	}
	if completionResponse.Error != nil {
		http.Error(w, "OpenAI API title error: "+completionResponse.Error.Message, http.StatusInternalServerError)
		return
	}
	if len(completionResponse.Choices) == 0 {
		http.Error(w, "No title returned from OpenAI API", http.StatusInternalServerError)
		return
	}
	responsePayload := map[string]string{
		"title": completionResponse.Choices[0].Message.Content,
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(responsePayload)
}

//
// File system endpoints
//

// filesListHandler lists all files (recursively) in the workspace.
func filesListHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Only GET allowed", http.StatusMethodNotAllowed)
		return
	}
	var files []string
	err := filepath.WalkDir(absWorkDir, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if !d.IsDir() {
			rel, err := filepath.Rel(absWorkDir, path)
			if err != nil {
				return err
			}
			files = append(files, rel)
		}
		return nil
	})
	if err != nil {
		http.Error(w, "Error listing files: "+err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{"files": files})
}

// fileHandler supports file read (GET) and write (POST) operations.
func fileHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		handleFileRead(w, r)
	case http.MethodPost:
		handleFileWrite(w, r)
	default:
		http.Error(w, "Only GET and POST allowed", http.StatusMethodNotAllowed)
	}
}

func handleFileRead(w http.ResponseWriter, r *http.Request) {
	filename := r.URL.Query().Get("filename")
	if filename == "" {
		http.Error(w, "Missing filename", http.StatusBadRequest)
		return
	}
	fullPath, err := secureFilePath(filename)
	if err != nil {
		http.Error(w, "Invalid filename: "+err.Error(), http.StatusBadRequest)
		return
	}
	data, err := os.ReadFile(fullPath)
	if err != nil {
		http.Error(w, "Error reading file: "+err.Error(), http.StatusInternalServerError)
		return
	}
	encoded := base64.StdEncoding.EncodeToString(data)
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"content_b64": encoded})
}

func handleFileWrite(w http.ResponseWriter, r *http.Request) {
	var payload struct {
		Filename   string `json:"filename"`
		ContentB64 string `json:"content_b64"`
	}
	if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
		http.Error(w, "Invalid JSON: "+err.Error(), http.StatusBadRequest)
		return
	}
	if payload.Filename == "" || payload.ContentB64 == "" {
		http.Error(w, "Filename and content_b64 are required", http.StatusBadRequest)
		return
	}
	decoded, err := base64.StdEncoding.DecodeString(payload.ContentB64)
	if err != nil {
		http.Error(w, "Error decoding Base64 content: "+err.Error(), http.StatusBadRequest)
		return
	}
	fullPath, err := secureFilePath(payload.Filename)
	if err != nil {
		http.Error(w, "Invalid filename: "+err.Error(), http.StatusBadRequest)
		return
	}
	dir := filepath.Dir(fullPath)
	if err := os.MkdirAll(dir, os.ModePerm); err != nil {
		http.Error(w, "Error creating directory: "+err.Error(), http.StatusInternalServerError)
		return
	}
	if err := os.WriteFile(fullPath, decoded, 0644); err != nil {
		http.Error(w, "Error writing file: "+err.Error(), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "success"})
}

func secureFilePath(relPath string) (string, error) {
	relPath = filepath.ToSlash(relPath)
	if strings.Contains(relPath, "..") {
		return "", errors.New("relative paths outside the workspace are not allowed")
	}
	fullPath := filepath.Join(absWorkDir, relPath)
	absFullPath, err := filepath.Abs(fullPath)
	if err != nil {
		return "", err
	}
	if !strings.HasPrefix(absFullPath, absWorkDir) {
		return "", errors.New("attempt to access files outside the workspace")
	}
	return absFullPath, nil
}