import 'dotenv/config';
import express from 'express';
import { Settings, VectorStoreIndex, storageContextFromDefaults, Document } from 'llamaindex';
import { agent } from "@llamaindex/workflow";
import timeout from 'connect-timeout';
import { OpenAI, OpenAIEmbedding } from '@llamaindex/openai';
import { QdrantVectorStore } from '@llamaindex/qdrant';
import cors from 'cors';
import * as fs from 'fs';
import generateContent from "./index_vertexai.js";
import dJSON from 'dirty-json';

// Initialize Express app
const app = express();
app.use(express.json());
app.use(timeout('8000s'));
app.use(cors());


// Setup LLM
Settings.llm = new OpenAI({
    model: process.env.OPENAI_MODEL,
    apiKey: process.env.OPENAI_API_KEY,
});

Settings.embedModel = new OpenAIEmbedding({
    apiKey: process.env.OPENAI_API_KEY,
});

// Initialize Qdrant vector store for plants
let plantsVectorStore;

async function initializePlantsVectorStore() {
    plantsVectorStore = new QdrantVectorStore({
        url: 'http://localhost:6333',
        collectionName: 'plants-rag',
    });
}

await initializePlantsVectorStore()

const reader = (await VectorStoreIndex.fromVectorStore(plantsVectorStore)).queryTool({
    metadata: {
        name: 'plants_growth_book',
        description: 'Use this reader to get information about plants and how they grow',
    },
});

const plantsAnalyzeAgent = agent({
    tools: [
        reader
    ],
    systemPrompt: `
    You are a world-class botanical scientist and speculative physicist AI. Your mission is to creatively and scientifically predict the outcome of exposing a plant to various, often extreme, environmental conditions.

    You MUST strictly adhere to the following rules:
    1.  Your primary source of truth is the provided as the plants_growth_book reader. You must prioritize information from this source as the scientific foundation for your predictions.
    2.  You are encouraged to use your **own general knowledge** of botany, chemistry, and physics to supplement the reader. This is crucial for interpreting novel scenarios (e.g., unusual substances, complex environmental interactions) that are not explicitly defined.
    3.  Your answers must be exclusively in PARSABLE JSON FORMAT. Do not include any additional text, explanations, or comments outside the JSON object.
`
})

const jsonFixerAgent = agent({
    tools: [],
    systemPrompt: `
    You are a JSON formatting expert. Your task is to ensure that the provided JSON string is properly formatted and valid. If the input is not valid JSON, you must correct it without altering its meaning.
    
    Out of whatever you receive, you must return a JSON object of this format:
    {
        "final_description_text": "...",
        "image_generator_prompt": "..."
    }
    
    DO NOT INCLUDE ABSOLUTELY ANYTHING EXCEPT FOR THE RAW JSON OBJECT. It needs to be parsed.
    `
})

const plantsHealerAgent = agent({
    tools: [
        reader
    ],
    systemPrompt: `
    You are a world-class botanical scientist and speculative physicist AI. Your mission is to provide advice on how to heal a plant that has been exposed to extreme environmental conditions.

    You MUST strictly adhere to the following rules:
    1.  Your primary source of truth is the provided as the plants_growth_book reader. You must prioritize information from this source as the scientific foundation for your predictions.
    2.  You are encouraged to use your **own general knowledge** of botany, chemistry, and physics to supplement the reader. This is crucial for interpreting novel scenarios (e.g., unusual substances, complex environmental interactions) that are not explicitly defined.
`
})

/**
 * Extracts the first complete JSON object from a string that might contain multiple objects
 * or other text.
 * @param {string} str The raw string from the API.
 * @returns {string|null} The cleaned string of the first JSON object, or null if none found.
 */
function extractFirstJsonObject(str) {
    let braceCount = 0;
    let startIndex = -1;

    for (let i = 0; i < str.length; i++) {
        if (str[i] === '{') {
            if (braceCount === 0) {
                startIndex = i;
            }
            braceCount++;
        } else if (str[i] === '}') {
            braceCount--;
            if (braceCount === 0 && startIndex !== -1) {
                // We found the matching closing brace for the first opening brace.
                return str.substring(startIndex, i + 1);
            }
        }
    }
    return null; // Return null if no complete object was found
}

app.post('/analyze', async (req, res) => {
    const query = `
        ### REASONING FRAMEWORK
        You must follow these steps to analyze the input:

        1.  **Deconstruct the Input:** Break down the "Environmental Changes & Actions" into their fundamental physical and chemical principles.
            *   **Example 1:** If the action is "watering with Coca-Cola," you must identify its key components: high sugar content (sucrose/fructose), high acidity (phosphoric and carbonic acid), and carbonation.
            *   **Example 2:** If the action is "nearby meteorite explosion," you must identify the immediate physical effects: intense thermal radiation (extreme heat), a concussive shockwave (extreme pressure), and potential ionizing radiation.

        2.  **Analyze the Causal Chain:** Identify the chain of **first-order, second-order, and third-order effects** on the plant.
            *   **First-order effects** are the immediate environmental changes (e.g., soil pH drops, air temperature skyrockets).
            *   **Second-order effects** are the direct impact on the plant's systems (e.g., roots can't absorb nutrients, cells are flash-boiled).
            *   **Third-order effects** are the final, observable morphological state (e.g., leaves are wilted and brown, the plant is vaporized).

        3.  **Synthesize and Quantify:** Combine all effects and factor in the "Time Period" to determine the final state and severity of the outcome. A short duration may only cause stress, while a long one could be lethal.
        ---
        ### INPUT DATA
        
        1. Initial Plant State:
            Description: ${JSON.stringify(req.body.plant_state_description)}
            
            Details: ${JSON.stringify(req.body.plant_state_details)}
        
        2. Environmental Changes & Actions:
            Description: ${JSON.stringify(req.body.environmental_changes_description)}
        
            Details: ${JSON.stringify(req.body.environmental_changes_details)}
        
        3. Time Period: ${JSON.stringify(req.body.time_period)}
        ---
        ### OUTPUT INSTRUCTIONS

        Based on your analysis using the REASONING FRAMEWORK, generate a single JSON object with two keys:
        
        1. "final_description_text": A detailed but clear textual description of the final state. Explain the "why" by referencing the scientific principles you identified (e.g., "Due to extreme heat from the thermal pulse..."). This should be 2-3 sentences long.
        2. "image_generator_prompt": "This MUST be a set of instructions for an AI image editor, telling it how to modify the initial photo to achieve the final state. It must be a comma-separated list of concise, actionable commands. Focus on the transformation itself, not the final picture. Use imperative verbs: 'Make...', 'Turn...', 'Add...', 'Reduce...', 'Introduce...', 'Bend...', 'Change...' Example: 'Turn the green leaves to a pale yellow, add wilting to the edges of all leaves, reduce the plant's overall bushiness by 30%, make the main stems thinner and longer, introduce small brown spots on some leaves'. Provide 5-10 clear, actionable commands."        
        
        Here's an example of the expected output format:
        {
            "final_description_text": "...",
            "image_generator_prompt": "..."
        }
        
        DO NOT INCLUDE ABSOLUTELY ANYTHING EXCEPT FOR THE RAW JSON OBJECT. It needs to be parsable. DO NOT repeat yourself.
        `

    const resultString = await runQuery(query, plantsAnalyzeAgent)

    console.log('🔍 Raw String Result:', resultString);
    // --- NEW, MORE ROBUST CLEANING LOGIC ---
    const jsonToParse = extractFirstJsonObject(resultString);

    if (!jsonToParse) {
        return res.status(500).json({
            success: false,
            error: "Could not extract any valid JSON object from the response.",
            rawResponse: resultString
        });
    }
    console.log('🔧 Extracted String for Parsing:', jsonToParse);

    const parsedResult = tryParseJSON(resultString);
    // BEST PRACTICE: Handle cases where the AI returns invalid JSON
    if (!parsedResult) {
        return res.status(500).json({
            success: false,
            error: "Failed to parse the extracted JSON object from the analysis service.",
            rawResponse: resultString
        });
    }
    res.json({
        success: true,
        result: parsedResult,
    });
})

app.post('/generate', async (req, res) => {
    const imageDescription = req.body[0];
    const imageBase64 = req.body[1];

    console.log('Image description:', imageDescription, 'Image base64 length:', imageBase64.length);

    const base64Data = await generateContent(imageBase64, imageDescription)

    res.json({
        success: true,
        result: `data:image/jpeg;base64,${base64Data}`,
    });
})

app.post('/cure', async (req, res) => {
    const query = `
       
        ### INPUT DATA
        
        1. Initial Plant State:
            Description: ${JSON.stringify(req.body.plant_state_description)}
            
            Details: ${JSON.stringify(req.body.plant_state_details)}
            
            Health Breakdown: ${JSON.stringify(req.body.health_breakdown)}
       
        ---
        
        Based on the input data and your reader, write 4-5 sentences of advice on how to heal the plant. End every sentence with a period.
        `

    const resultString = await runQuery(query, plantsHealerAgent)

    res.json({
        success: true,
        result: resultString,
    });
})


// Helper function to run queries
async function runQuery(query, ragAgent) {
    const response = await ragAgent.run(query);

    if (process.env.LOGGING_LEVEL === '1') {
        console.log('✅ Query result:', response.data.result);
    }

    return response.data.result;
}

// Helper to safely parse JSON
async function tryParsedJSON(jsonString) {
    try {
        return dJSON.parse(jsonString);
    } catch (e) {
        console.error('❌ JSON parsing error:', e.message);
        const result = await runQuery(`Fix this JSON: ${jsonString}`, jsonFixerAgent)

        console.log('🔧 Fixed JSON:', result);
        return dJSON.parse(result);
    }
}

// A robust function to parse JSON safely
function tryParseJSON(jsonString) {
    try {
        const o = JSON.parse(jsonString);
        // Handle non-exception-throwing cases:
        // JSON.parse(null) returns null, and other edge cases.
        if (o && typeof o === "object") {
            return o;
        }
    } catch (e) {
        console.error("Error parsing JSON string:", e);
        console.error("Original string was:", jsonString);
    }
    return null; // Return null if parsing fails
}

// Start the server
const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {
    console.log(`Server running on port ${PORT}`);
});
