package main

import (
	"encoding/json"
	"fmt"
	"inou/lib"
	"os"
	"path/filepath"
	"strings"
)

// --- Local Structs for Prompt Processing ---
// These are defined locally to this application as per the new architecture.

type TriageResponse struct {
	Category string `json:"category"`
	Language string `json:"language"`
	HasData  bool   `json:"has_data"`
	Error    string `json:"error,omitempty"`
}

type ExtractionResult struct {
	Question    string         `json:"question"`
	Category    string         `json:"category"`
	Type        string         `json:"type"`
	InputType   string         `json:"input_type"`
	InputConfig InputConfig    `json:"input_config"`
	Schedule    []ScheduleSlot `json:"schedule"`
	Entries     []*EntryData   `json:"entries,omitempty"`
	Error       string         `json:"error,omitempty"`
}

type InputConfig struct {
	Fields []FormField `json:"fields,omitempty"`
	Groups []FormGroup `json:"groups,omitempty"`
}

type FormGroup struct {
	Title  string      `json:"title"`
	Fields []FormField `json:"fields"`
}

type FormField struct {
	Key     string   `json:"key"`
	Type    string   `json:"type"`
	Label   string   `json:"label"`
	Unit    string   `json:"unit,omitempty"`
	Options []string `json:"options,omitempty"`
}

type ScheduleSlot struct {
	Days []string `json:"days"`
	Time string   `json:"time"`
}

type EntryData struct {
	Value string      `json:"value"`
	Data  interface{} `json:"data"`
}

var ValidCategories = map[string]bool{
	"vital": true, "exercise": true, "medication": true, "supplement": true,
	"symptom": true, "note": true, "surgery": true, "hospitalization": true,
	"consultation": true, "diagnosis": true, "device": true, "therapy": true,
	"assessment": true, "birth": true, "imaging_finding": true, "eeg_finding": true,
	"provider": true, "question": true, "history": true, "family_history": true,
	"nutrition": true, "fertility": true, "out_of_domain": true,
}


// --- Main Application Logic ---

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: test-prompts <input>")
		os.Exit(1)
	}
	input := strings.Join(os.Args[1:], " ")

	// Initialize shared library components
	lib.ConfigInit() // Sets lib.GeminiKey from .env
	lib.InitPrompts("api/prompts") // Sets the prompts dir for the lib

	// Triage
	fmt.Println("=== TRIAGE ===")
	triage, err := runTriage(input)
	if err != nil {
		fmt.Printf("Triage error: %v\n", err)
		os.Exit(1)
	}
	if triage.Error != "" {
		fmt.Printf("Error: %s\n", triage.Error)
		os.Exit(0)
	}
	fmt.Printf("Category: %s\nLanguage: %s\nHas data: %v\n", triage.Category, triage.Language, triage.HasData)
	

	// Extraction
	fmt.Println("\n=== EXTRACTION ===")
	// For a standalone test, we don't have existing types to pass.
	existingTypes := make(map[string][]string)

	// Special case for testing pregnancy contradiction
	if input == "I just had my period" {
		existingTypes["fertility"] = []string{"pregnancy"}
		fmt.Println("INFO: Injected existing type: fertility: [pregnancy]")
	}

	result, err := runExtraction(input, triage.Category, triage.Language, existingTypes)
	if err != nil {
		fmt.Printf("Extraction error: %v\n", err)
		os.Exit(1)
	}

	out, _ := json.MarshalIndent(result, "", "  ")
	fmt.Println(string(out))

	// Summary
	fmt.Println("\n=== SUMMARY ===")
	if len(result.Entries) > 0 {
		for i, entry := range result.Entries {
			fmt.Printf("Entry %d: %s\n", i+1, entry.Value)
		}
	} else {
		fmt.Println("Entry: (none)")
	}

	if len(result.Schedule) > 0 && result.Question != "" {
		fmt.Printf("Prompt: %s\n", result.Question)
		fmt.Printf("Schedule contains %d slot(s)\n", len(result.Schedule))
	} else {
		fmt.Println("Prompt: (none - one-time event)")
	}
}


// --- Local Prompt Handling Functions ---

func loadPrompt(name string) (string, error) {
	path := filepath.Join(lib.TrackerPromptsDir(), name+".md")
	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

func runTriage(userInput string) (*TriageResponse, error) {
	tmpl, err := loadPrompt("triage")
	if err != nil {
		return nil, fmt.Errorf("failed to load triage prompt: %v", err)
	}

	// Use same context as extraction
	dossierName := "Sophia Helena"
	dossierDOB := "2017-01-01"

	prompt := strings.ReplaceAll(tmpl, "{{INPUT}}", userInput)
	prompt = strings.ReplaceAll(prompt, "{{DOSSIER_NAME}}", dossierName)
	prompt = strings.ReplaceAll(prompt, "{{DOSSIER_DOB}}", dossierDOB)
	
	respText, err := lib.CallGemini(prompt)
	if err != nil {
		return nil, err
	}

	var result TriageResponse
	if err := json.Unmarshal([]byte(respText), &result); err != nil {
		var errMap map[string]string
		if json.Unmarshal([]byte(respText), &errMap) == nil {
			if errMsg, ok := errMap["error"]; ok {
				result.Error = errMsg
				return &result, nil
			}
		}
		return nil, fmt.Errorf("failed to parse triage JSON: %v (raw: %s)", err, respText)
	}

	if _, ok := ValidCategories[result.Category]; !ok && result.Error == "" {
		result.Category = "note"
	}
	return &result, nil
}

func runExtraction(userInput, category, language string, existingTypes map[string][]string) (*ExtractionResult, error) {
	tmpl, err := loadPrompt(category)
	if err != nil {
		tmpl, err = loadPrompt("default")
		if err != nil {
			return nil, fmt.Errorf("failed to load prompt: %v", err)
		}
	}

	var existingStr string
	for cat, types := range existingTypes {
		if len(types) > 0 {
			existingStr += fmt.Sprintf("- %s: %v\n", cat, types)
		}
	}
	if existingStr == "" {
		existingStr = "(none yet)"
	}

	// Use Sophia's context for testing (child patient helps with practitioner detection)
	dossierName := "Sophia Helena"
	dossierDOB := "2017-01-01"
	currentDate := "2026-02-08"

	prompt := tmpl
	prompt = strings.ReplaceAll(prompt, "{{INPUT}}", userInput)
	prompt = strings.ReplaceAll(prompt, "{{LANGUAGE}}", language)
	prompt = strings.ReplaceAll(prompt, "{{CATEGORY}}", category)
	prompt = strings.ReplaceAll(prompt, "{{EXISTING_TYPES}}", existingStr)
	prompt = strings.ReplaceAll(prompt, "{{DOSSIER_NAME}}", dossierName)
	prompt = strings.ReplaceAll(prompt, "{{DOSSIER_DOB}}", dossierDOB)
	prompt = strings.ReplaceAll(prompt, "{{CURRENT_DATE}}", currentDate)
	
	respText, err := lib.CallGemini(prompt)
	if err != nil {
		return nil, err
	}

	var result ExtractionResult
	if err := json.Unmarshal([]byte(respText), &result); err != nil {
		// Fallback for single entry to maintain compatibility with older prompts
		var singleEntryResult struct {
			Question    string      `json:"question"`
			Category    string      `json:"category"`
			Type        string      `json:"type"`
			InputType   string      `json:"input_type"`
			InputConfig InputConfig `json:"input_config"`
			Schedule    []ScheduleSlot `json:"schedule"`
			Entry       *EntryData     `json:"entry,omitempty"`
			Error       string      `json:"error,omitempty"`
		}
		if err2 := json.Unmarshal([]byte(respText), &singleEntryResult); err2 == nil {
			result.Question = singleEntryResult.Question
			result.Category = singleEntryResult.Category
			result.Type = singleEntryResult.Type
			result.InputType = singleEntryResult.InputType
			result.InputConfig = singleEntryResult.InputConfig
			result.Schedule = singleEntryResult.Schedule
			result.Error = singleEntryResult.Error
			if singleEntryResult.Entry != nil {
				result.Entries = []*EntryData{singleEntryResult.Entry}
			}
			if result.Category == "" {
				result.Category = category
			}
			return &result, nil
		}
		return nil, fmt.Errorf("failed to parse extraction JSON: %v (raw: %s)", err, respText)
	}

	if result.Category == "" {
		result.Category = category
	}
	return &result, nil
}