package main

import (
	"bytes"
	"encoding/json"
	"html/template"
	"io"
	"log"
	"net/http"
	"os"
	"time"
)

// Message represents a single 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"`
	// no temperature field for "o3-mini"
}

// ChatCompletionResponseChoice represents one of the returned choices.
type ChatCompletionResponseChoice struct {
	Message Message `json:"message"`
}

// ChatCompletionResponse is the API response structure.
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
)

func init() {
	openaiAPIKey = os.Getenv("OPENAI_API_KEY")
	if openaiAPIKey == "" {
		log.Fatal("Environment variable OPENAI_API_KEY is not set")
	}

	indexTmpl = template.Must(template.New("index").Parse(indexHTML))
}

func main() {
	http.HandleFunc("/", indexHandler)
	http.HandleFunc("/chat", chatHandler)
	http.HandleFunc("/title", titleHandler) // New endpoint for titling conversation

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

// indexHandler serves the main HTML page.
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)
	}
}

// chatHandler receives conversation messages from the client,
// calls the OpenAI API, and returns the assistant's reply.
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)
}

// titleHandler provides an automatic title for the conversation.
// It expects a POST with JSON: { messages: []Message } and then
// calls the OpenAI API with a special prompt to generate a title.
func titleHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Only POST allowed", http.StatusMethodNotAllowed)
		return
	}

	// Read the incoming JSON containing the conversation messages.
	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
	}

	// Prepend a system message that instructs the AI to generate
	// a short title based on the conversation.
	titleRequestMessages := []Message{
		{
			Role:    "system",
			Content: "You are an AI that generates conversation titles. Given the conversation below, provide a short, descriptive title (in no more than five words) with no additional 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)
}

const indexHTML = `
<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>ChatGPT Clone - Conversations</title>
  <meta name="viewport" content="width=device-width, initial-scale=1">
  <!-- Bootstrap CSS -->
  <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
  <!-- Highlight.js CSS -->
  <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/styles/default.min.css">
  <style>
    body { background-color: #f7f7f7; }
    .container-main { margin-top: 1rem; }
    .sidebar {
      max-height: 80vh; overflow-y: auto; background-color: #fff;
      border-radius: 8px; padding: 0.5rem; box-shadow: 0 0 10px rgba(0,0,0,0.1);
    }
    .chat-container {
      background-color: #fff; border-radius: 8px; padding: 1rem;
      box-shadow: 0 0 10px rgba(0,0,0,0.1); flex-grow: 1;
    }
    .message {
      padding: 0.75rem; border-radius: 5px; margin-bottom: 0.5rem; white-space: pre-wrap;
    }
    .user { background-color: #d1e7dd; text-align: right; }
    .assistant { background-color: #f8d7da; text-align: left; }
    .chat-log {
      max-height: 400px; overflow-y: auto; margin-bottom: 1rem;
      background-color: #eee; border: 1px solid #ddd; border-radius: 5px; padding: 0.5rem;
    }
    pre {
      background-color: #2d2d2d; color: #f8f8f2; padding: 0.5rem;
      border-radius: 5px; overflow-x: auto;
    }
    .conversation-item {
      padding: 0.5rem; cursor: pointer; border-bottom: 1px solid #ddd; position: relative;
    }
    .conversation-item:hover, .conversation-item.active { background-color: #f0f0f0; }
    /* Delete button styling */
    .delete-btn {
      position: absolute; right: 5px; top: 5px; background: transparent;
      border: none; color: #dc3545; font-size: 1rem;
    }
    /* Style the textarea for multi-line messages */
    textarea#messageInput {
      resize: none; height: 80px; overflow-y: auto;
    }
  </style>
</head>
<body>
  <div class="container container-main">
    <div class="row">
      <!-- Sidebar: Conversation List -->
      <div class="col-md-3">
        <div class="sidebar">
          <h5>Conversations</h5>
          <button id="newConvoBtn" class="btn btn-sm btn-primary mb-2">New Conversation</button>
          <div id="conversationList"></div>
        </div>
      </div>
      <!-- Chat Container -->
      <div class="col-md-9">
        <div class="chat-container d-flex flex-column">
          <h4 id="chatTitle">Conversation</h4>
          <div id="chatLog" class="chat-log"></div>
          <!-- Changed input to textarea for multi-line messages -->
          <form id="chatForm" class="mt-auto">
            <div class="input-group">
              <textarea id="messageInput" class="form-control" placeholder="Type your message here" required></textarea>
              <button class="btn btn-primary" type="submit">Send</button>
            </div>
          </form>
        </div>
      </div>
    </div>
  </div>

  <!-- Marked & Highlight.js Libraries -->
  <script src="https://cdnjs.cloudflare.com/ajax/libs/marked/4.3.0/marked.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.7.0/highlight.min.js"></script>
  <script>
    // Configure marked to use highlight.js for code blocks.
    marked.setOptions({
      highlight: function(code, lang) {
        if (lang && hljs.getLanguage(lang)) {
          return hljs.highlight(code, { language: lang }).value;
        }
        return hljs.highlightAuto(code).value;
      }
    });

    // --- Conversation Management ---
    // Each conversation is an object: { id: string, title: string, messages: [ {role, content} ] }
    // Stored in localStorage under key "conversations"
    let conversations = [];
    let currentConversation = null;

    // Generate unique conversation IDs.
    function generateId() {
      return 'c-' + Date.now() + '-' + Math.floor(Math.random()*1000);
    }

    function loadConversations() {
      const loaded = localStorage.getItem('conversations');
      if (loaded) {
        conversations = JSON.parse(loaded);
      } else {
        conversations = [];
      }
      renderConversationList();
    }

    function saveConversations() {
      localStorage.setItem('conversations', JSON.stringify(conversations));
    }

    function createNewConversation() {
      const newConvo = {
        id: generateId(),
        title: 'New Conversation',
        messages: [
          { role: "system", content: "You are ChatGPT, a helpful assistant. When providing code, please always wrap the code in markdown triple backticks so that it renders correctly." }
]
};
conversations.push(newConvo);
currentConversation = newConvo;
saveConversations();
renderConversationList();
renderChatLog();
document.getElementById('chatTitle').textContent = "Conversation (" + newConvo.id + ")";
}

function deleteConversation(id) {
if (!confirm("Are you sure you want to delete this conversation?")) return;
conversations = conversations.filter(convo => convo.id !== id);
if (currentConversation && currentConversation.id === id) {
currentConversation = conversations.length > 0 ? conversations[0] : null;
if (!currentConversation) { createNewConversation(); }
else { document.getElementById('chatTitle').textContent = "Conversation (" + currentConversation.id + ")"; }
}
saveConversations();
renderConversationList();
renderChatLog();
}

function renderConversationList() {
const listElem = document.getElementById('conversationList');
listElem.innerHTML = "";
conversations.forEach(convo => {
const div = document.createElement("div");
div.className = "conversation-item" + (currentConversation && currentConversation.id === convo.id ? " active" : "");
div.textContent = convo.title || convo.id;

// Create delete button.
const deleteBtn = document.createElement("button");
deleteBtn.className = "delete-btn";
deleteBtn.innerHTML = "&times;";
deleteBtn.onclick = function(e) { e.stopPropagation(); deleteConversation(convo.id); };
div.appendChild(deleteBtn);

// Click on conversation selects it.
div.onclick = () => {
currentConversation = convo;
document.getElementById('chatTitle').textContent = "Conversation (" + convo.id + ")";
renderConversationList();
renderChatLog();
};
listElem.appendChild(div);
});
}

function renderChatLog() {
const chatLog = document.getElementById('chatLog');
chatLog.innerHTML = "";
if (!currentConversation) return;
currentConversation.messages.forEach(msg => {
appendMessage(msg.role, msg.content, false);
});
chatLog.scrollTop = chatLog.scrollHeight;
}

// Append message to chat log.
// "update" determines if the message is added to the conversation.
function appendMessage(role, content, update = true) {
const chatLog = document.getElementById('chatLog');
const messageElem = document.createElement("div");
messageElem.className = "message " + role;
messageElem.innerHTML = marked.parse(content);
chatLog.appendChild(messageElem);
chatLog.scrollTop = chatLog.scrollHeight;
if (update && currentConversation) {
currentConversation.messages.push({ role: role, content: content });
// Auto-update title if still "New Conversation" after the assistant's reply.
if (role === "assistant" && currentConversation.title === "New Conversation") {
autoTitleConversation();
}
saveConversations();
}
return messageElem;
}

// Call the /title endpoint to auto-generate a conversation title,
// then update the conversation and UI.
async function autoTitleConversation() {
try {
const response = await fetch("/title", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ messages: currentConversation.messages })
});
if (!response.ok) {
throw new Error("Title API error: " + response.status);
}
const data = await response.json();
// Use the returned title (trim it) to update the conversation.
const newTitle = data.title.trim();
if(newTitle) {
currentConversation.title = newTitle;
document.getElementById('chatTitle').textContent = "Conversation (" + currentConversation.id + "): " + newTitle;
renderConversationList();
saveConversations();
}
} catch (error) {
console.error("Error auto-titling conversation:", error);
}
}

// --- Chat Form Submission ---
const chatForm = document.getElementById('chatForm');
const messageInput = document.getElementById('messageInput');

// Allow Shift+Enter for newlines.
messageInput.addEventListener('keydown', function(e) {
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
chatForm.requestSubmit();
}
});

chatForm.addEventListener('submit', async function(e) {
e.preventDefault();
const userMessage = messageInput.value.trim();
if (!userMessage || !currentConversation) return;

appendMessage("user", userMessage);
messageInput.value = "";

// Add temporary typing indicator.
const typingIndicator = appendMessage("assistant", "Typing...", false);

try {
const response = await fetch("/chat", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ messages: currentConversation.messages })
});
if (!response.ok) {
throw new Error("Server error: " + response.status);
}
const data = await response.json();
typingIndicator.remove();
appendMessage("assistant", data.reply);
} catch (error) {
typingIndicator.remove();
appendMessage("assistant", "Error: " + error.message);
}
});

// New Conversation button.
document.getElementById('newConvoBtn').addEventListener("click", function() {
createNewConversation();
});

// --- Initialization ---
loadConversations();
if (conversations.length === 0) {
createNewConversation();
} else {
currentConversation = conversations[0];
document.getElementById('chatTitle').textContent = "Conversation (" + currentConversation.id + ")";
renderChatLog();
}
</script>
</body>
</html>
`