main.go

package main

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"regexp"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
	_ "github.com/lib/pq"
)

// ==========================================
// 1. CONFIGURATION
// ==========================================
const (
	BrainURL = "https://triflix-brainfuncall.hf.space/generate"
	DB_DSN   = "postgres://avnadmin:AVNS_1YpuKQZB2rqpqAh8p9o@pravah-ai-project-e131.h.aivencloud.com:25145/defaultdb?sslmode=require"
	Port     = ":7860"
)

const CA_CERT = `-----BEGIN CERTIFICATE-----
MIIEUDCCArigAwIBAgIUYnE2Ndy3Mu5ZfkaYqb9wNwepN0UwDQYJKoZIhvcNAQEM
BQAwQDE+MDwGA1UEAww1M2JmOTUzODgtMzM3Yi00MWQ5LTgwYzEtMWMyY2MyZWE5
N2U2IEdFTiAxIFByb2plY3QgQ0EwHhcNMjUxMjE2MDM1MDI4WhcNMzUxMjE0MDM1
MDI4WjBAMT4wPAYDVQQDDDUzYmY5NTM4OC0zMzdiLTQxZDktODBjMS0xYzJjYzJl
YTk3ZTYgR0VOIDEgUHJvamVjdCBDQTCCAaIwDQYJKoZIhvcNAQEBBQADggGPADCC
AYoCggGBAOstPCPmXXsVkGCGb4y/6coATrk7OxmLVpw0X0yq9JQ6JJif+PcRTQ66
kV0P7Kf88Zip0z/MNeAbar5j3f01EhEO+QsQFUpbZw9BkVxyt6g1T8NPDq3qRBtr
ArYqKFWCSv8IrjRlBz/VIQ7NrbwljyDciXKPgXwK98loEmQAzDnD5d3L5MvH/YKN
r1LRsXiqtwmGox6/EQTUr2MvR4vpIg70HWwdIQOf2/nA97moMrmu/dBx58VyBadl
rwWCFn/efb5fwlRlDzYcc8unPDA4luRAOr7EKJC7Aj3z45CynH+NJWxCeCagL1fA
b/apcO4cVOge0xTn1U1Kpj6ujmqoL7ZW+kdBj6qlYRKPUbdRXE6gxNWtyjHRXvCm
hFeTivS8twLtJxDigi3Gv8ynsikFnSk3SrK61ptWEgt9y/KmEjB1qJ1RUoRBJDB0
JSMxKN+d4CrEbhl5zZ87jN8DTqcp86AsgP3X12JJltSa/s5SaU7DJPOOKP3hcHdP
KnyXquqITwIDAQABo0IwQDAdBgNVHQ4EFgQUcSCbtk4vb0mcmcIFlg1AWi1JHnEw
EgYDVR0TAQH/BAgwBgEB/wIBADALBgNVHQ8EBAMCAQYwDQYJKoZIhvcNAQEMBQAD
ggGBAAkbXC6kli1nk1eZW5YBe6ObQeow5FVE5LC6Ft9nW+hCUvYcGQpvEAuyeKpc
i7wmo914WL0zJbQLmU5l2U5TCy+dpu34VqJ87c0CGO9z2vXQavXpz7l02QnEL77I
4sWx4lVmdgqxdq0GpQHyhiCTCDPCEPcKX3jOcEdgRKv0V8N29cA/WsS/AOS5fqoc
RAjDllxyMYbByhrXSPuDUfW5dKABfMvmBsIEqNki8BejYmuMc+CE+i/lOwWQ4D7U
smw+OZEE5I9aMnDN9A7jA/ArCf8bJI57iEspsu4vOIYQgqVw21Br37afp18CePxF
yqHr4EaHQMFFd1jH03nvSiRQREd1bfjB50+mongOdltRHm4yBQ5IthfwMlwXvDrX
/YHU2XLLg3XnsUwd951FL6dx5FXgYPPWEGiifWwreZggjRDkp+UfxhJCo8fKPzNS
oKgr4gM6XleFIcnlW1vJqxoTuybLmgJ9GM/L++gpBd4Fk5gZtZOVqY0iiHr1zI3o
gnWFGw==
-----END CERTIFICATE-----`

// ==========================================
// 2. ADVANCED TOOLS SCHEMA
// ==========================================
var ToolsSchema = []map[string]interface{}{
	// TOOL 1: Technical Data (Water Levels, Storage)
	{
		"type": "function",
		"function": map[string]interface{}{
			"name":        "search_technical_data",
			"description": "Get water level, storage TMC, and percentage for a specific dam.",
			"parameters": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dam_name": map[string]interface{}{"type": "string"},
				},
				"required": []string{"dam_name"},
			},
		},
	},
	// TOOL 2: Admin Data (Corporation, Division)
	{
		"type": "function",
		"function": map[string]interface{}{
			"name":        "search_admin_details",
			"description": "Find who manages a dam (Corporation, Division) or who entered the data.",
			"parameters": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"dam_name": map[string]interface{}{"type": "string"},
				},
				"required": []string{"dam_name"},
			},
		},
	},
	// TOOL 3: Advanced Filtering (Taluka, Purpose, Type)
	{
		"type": "function",
		"function": map[string]interface{}{
			"name":        "filter_dams_advanced",
			"description": "Search dams by Location (Taluka/District), Purpose (Irrigation), or Type.",
			"parameters": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"district":       map[string]interface{}{"type": "string"},
					"taluka":         map[string]interface{}{"type": "string"},
					"purpose":        map[string]interface{}{"type": "string", "description": "e.g. Irrigation, Hydroelectricity"},
					"project_type":   map[string]interface{}{"type": "string", "enum": []string{"Major", "Medium", "Minor"}},
					"min_percentage": map[string]interface{}{"type": "integer"},
				},
			},
		},
	},
	// TOOL 4: Regional Statistics
	{
		"type": "function",
		"function": map[string]interface{}{
			"name":        "get_region_statistics",
			"description": "Get the Average Water Storage (%) for a whole District or Region.",
			"parameters": map[string]interface{}{
				"type": "object",
				"properties": map[string]interface{}{
					"location_name": map[string]interface{}{"type": "string", "description": "Name of District or Revenue Region"},
				},
				"required": []string{"location_name"},
			},
		},
	},
}

// ==========================================
// 3. DATA STRUCTURES
// ==========================================
type ChatRequest struct {
	Message string `json:"message"`
}

type BrainPayload struct {
	Query       string                   `json:"query"`
	Tools       []map[string]interface{} `json:"tools"`
	IncludeDate bool                     `json:"include_date"`
}

type BrainResponse struct {
	Response string `json:"response"`
}

// ==========================================
// 4. DATABASE INIT
// ==========================================
var db *sql.DB

func initDB() {
	tmpFile, err := ioutil.TempFile("", "ca-*.pem")
	if err != nil {
		log.Fatal("Failed to create temp cert file:", err)
	}
	if _, err := tmpFile.Write([]byte(CA_CERT)); err != nil {
		log.Fatal("Failed to write cert:", err)
	}
	tmpFile.Close()

	connStr := fmt.Sprintf("%s&sslrootcert=%s", DB_DSN, tmpFile.Name())
	var dbErr error
	db, dbErr = sql.Open("postgres", connStr)
	if dbErr != nil {
		log.Fatal("Failed to open DB:", dbErr)
	}
	if err := db.Ping(); err != nil {
		log.Fatal("Failed to ping DB:", err)
	}
	fmt.Println("✅ Connected to Aiven PostgreSQL")
}

// ==========================================
// 5. MAIN SERVER
// ==========================================
func main() {
	initDB()
	gin.SetMode(gin.ReleaseMode)
	r := gin.Default()

	r.GET("/", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "Golang Muscle is Active"})
	})

	r.POST("/chat", func(c *gin.Context) {
		var req ChatRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(400, gin.H{"error": "Invalid JSON"})
			return
		}

		fmt.Printf("🗣️ User: %s\n", req.Message)

		// 1. Call Brain
		brainResp, err := callBrain(req.Message)
		if err != nil {
			fmt.Printf("❌ Brain Error: %v\n", err)
			c.JSON(500, gin.H{"error": "Brain is offline"})
			return
		}
		fmt.Printf("🤖 Brain Said: %s\n", brainResp)

		// 2. Route Logic (Priority System)
		finalAnswer := "I'm not sure how to answer that."

		if strings.Contains(brainResp, "call:search_technical_data") {
			dName := extractArg(brainResp, "dam_name")
			finalAnswer = executeTechnicalSearch(dName)
		
		} else if strings.Contains(brainResp, "call:search_admin_details") {
			dName := extractArg(brainResp, "dam_name")
			finalAnswer = executeAdminSearch(dName)

		} else if strings.Contains(brainResp, "call:get_region_statistics") {
			loc := extractArg(brainResp, "location_name")
			finalAnswer = executeRegionStats(loc)

		} else if strings.Contains(brainResp, "call:filter_dams_advanced") {
			dist := extractArg(brainResp, "district")
			taluka := extractArg(brainResp, "taluka")
			purpose := extractArg(brainResp, "purpose")
			pType := extractArg(brainResp, "project_type")
			minPct := extractArg(brainResp, "min_percentage")
			finalAnswer = executeAdvancedFilter(dist, taluka, purpose, pType, minPct)
		
		} else {
			finalAnswer = "⚠️ I couldn't find a tool to answer that."
		}

		c.JSON(200, gin.H{
			"answer":    finalAnswer,
			"raw_brain": brainResp,
		})
	})

	fmt.Printf("🚀 Server listening on %s\n", Port)
	r.Run(Port)
}

// ==========================================
// 6. BRAIN COMMUNICATION
// ==========================================
func callBrain(query string) (string, error) {
	payload := BrainPayload{
		Query:       query, 
		Tools:       ToolsSchema,
		IncludeDate: true,
	}

	jsonData, _ := json.Marshal(payload)
	resp, err := http.Post(BrainURL, "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var result BrainResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}
	return result.Response, nil
}

// ==========================================
// 7. ROBUST PARSER
// ==========================================
func extractArg(text, key string) string {
	cleanText := strings.ReplaceAll(text, "<escape>", "")
	re := regexp.MustCompile(key + `:([^,^}]+)`)
	matches := re.FindStringSubmatch(cleanText)
	if len(matches) > 1 {
		return strings.TrimSpace(matches[1])
	}
	return ""
}

// ==========================================
// 8. SQL EXECUTORS (The Muscle)
// ==========================================

// Tool 1: Technical Data
func executeTechnicalSearch(damName string) string {
	cleanName := strings.ToLower(damName)
	cleanName = strings.ReplaceAll(cleanName, " dam", "")
	
	query := `
		SELECT name_of_dam, district, lake_level_reading_rl, live_storage_tmc, live_storage_designed_live_storage__pct, snapshot_report_date
		FROM dam_water_levels 
		WHERE name_of_dam ILIKE $1 
		LIMIT 1`

	var name, dist, date string
	var level, tmc, pct float64

	err := db.QueryRow(query, "%"+cleanName+"%").Scan(&name, &dist, &level, &tmc, &pct, &date)
	if err == sql.ErrNoRows {
		return fmt.Sprintf("❌ Could not find dam '%s'.", damName)
	}
	return fmt.Sprintf("🌊 **%s** (%s)\n   • Level: %.2f m\n   • Storage: %.3f TMC\n   • Filled: %.2f%%\n   • Date: %s", name, dist, level, tmc, pct, date)
}

// Tool 2: Admin Data
func executeAdminSearch(damName string) string {
	cleanName := strings.ToLower(damName)
	cleanName = strings.ReplaceAll(cleanName, " dam", "")

	query := `
		SELECT name_of_dam, corporation, division, data_entered_by
		FROM dam_water_levels 
		WHERE name_of_dam ILIKE $1 
		LIMIT 1`

	var name, corp, div, user string
	err := db.QueryRow(query, "%"+cleanName+"%").Scan(&name, &corp, &div, &user)
	if err == sql.ErrNoRows {
		return fmt.Sprintf("❌ Could not find dam '%s'.", damName)
	}
	return fmt.Sprintf("🏢 **%s**\n   • Corp: %s\n   • Division: %s\n   • Data Entry: %s", name, corp, div, user)
}

// Tool 3: Advanced Filter
func executeAdvancedFilter(district, taluka, purpose, pType, minPctStr string) string {
	baseQuery := "SELECT name_of_dam, district, taluka, live_storage_designed_live_storage__pct FROM dam_water_levels WHERE 1=1"
	var args []interface{}
	argCounter := 1

	if district != "" {
		baseQuery += fmt.Sprintf(" AND district ILIKE $%d", argCounter)
		args = append(args, "%"+district+"%")
		argCounter++
	}
	if taluka != "" {
		baseQuery += fmt.Sprintf(" AND taluka ILIKE $%d", argCounter)
		args = append(args, "%"+taluka+"%")
		argCounter++
	}
	if purpose != "" {
		baseQuery += fmt.Sprintf(" AND purpose_of_dam ILIKE $%d", argCounter)
		args = append(args, "%"+purpose+"%")
		argCounter++
	}
	if pType != "" {
		baseQuery += fmt.Sprintf(" AND project_type ILIKE $%d", argCounter)
		args = append(args, pType)
		argCounter++
	}
	if minPctStr != "" {
		minPct, _ := strconv.Atoi(minPctStr)
		baseQuery += fmt.Sprintf(" AND CAST(live_storage_designed_live_storage__pct AS DECIMAL) >= $%d", argCounter)
		args = append(args, minPct)
		argCounter++
	}

	baseQuery += " LIMIT 5"

	rows, err := db.Query(baseQuery, args...)
	if err != nil {
		return "DB Error: " + err.Error()
	}
	defer rows.Close()

	var results []string
	for rows.Next() {
		var name, dist, tal string
		var pct float64
		rows.Scan(&name, &dist, &tal, &pct)
		results = append(results, fmt.Sprintf("- %s (%s, %s): %.1f%%", name, dist, tal, pct))
	}

	if len(results) == 0 {
		return "No dams found matching those filters."
	}
	return "🔎 **Found Dams:**\n" + strings.Join(results, "\n")
}

// Tool 4: Region Stats
func executeRegionStats(location string) string {
	// Try to match District OR Region
	query := `
		SELECT COUNT(*), AVG(CAST(live_storage_designed_live_storage__pct AS DECIMAL))
		FROM dam_water_levels
		WHERE district ILIKE $1 OR revenue_region ILIKE $1`

	var count int
	var avgPct sql.NullFloat64 // Handle nulls if no data

	err := db.QueryRow(query, "%"+location+"%").Scan(&count, &avgPct)
	if err != nil {
		return "DB Error: " + err.Error()
	}

	if count == 0 {
		return fmt.Sprintf("❌ No data found for location '%s'.", location)
	}

	val := 0.0
	if avgPct.Valid {
		val = avgPct.Float64
	}

	return fmt.Sprintf("📊 **Statistics for %s**\n   • Total Dams: %d\n   • Average Storage: %.2f%%", location, count, val)
}