import streamlit as st
import json
import time
from datetime import datetime
from openai import OpenAI
from prompts import FERTILIZATION_PLAN_PROMPT_CORRECTED

# ===============================
# CONFIGURATION
# ===============================
OPENROUTER_API_KEY = "Add-your-openrouter-key"

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY,
)

# ===============================
# STREAMLIT UI
# ===============================
st.set_page_config(
    page_title="ThinkLawn Fertilization Plan Generator",
    page_icon="🌱",
    layout="centered",
)

st.title("🌿 ThinkLawn Fertilization Plan Generator")
st.caption("Generate professional fertilization plans based on soil test data and location details.")

# Input area
st.subheader("🧾 Enter Input Data (JSON Format or Text)")
user_input = st.text_area(
    "Paste your soil and questionnaire data below:",
    height=350,
    placeholder="Paste your JSON input here...",
)

# When user submits
if st.button("🚀 Generate Fertilization Plan"):
    if not user_input.strip():
        st.error("Please provide input data before submitting.")
    else:
        with st.spinner("Generating fertilization plan... Please wait ⏳"):
            start_time = time.time()

            try:
                # Try to load JSON (fallback to treat as text)
                try:
                    input_data = json.loads(user_input)
                    formatted_input = json.dumps(input_data, indent=2)
                except json.JSONDecodeError:
                    formatted_input = user_input.strip()

                user_message = f"Generate fertilization plan:\n\n{formatted_input}"

                # LLM Call
                response = client.chat.completions.create(
                    model="openai/gpt-4o-mini",
                    messages=[
                        {"role": "system", "content": FERTILIZATION_PLAN_PROMPT_CORRECTED},
                        {"role": "user", "content": user_message}
                    ],
                    temperature=0.3,
                    extra_headers={
                        "HTTP-Referer": "https://github.com/thinklawn/fertilization-optimizer",
                        "X-Title": "ThinkLawn Corrected Schema"
                    }
                )

                # Extract and clean response
                content = response.choices[0].message.content.strip()
                if content.startswith("```"):
                    content = content.split("```")[1]
                    if content.startswith("json"):
                        content = content[4:]
                    content = content.strip()

                # Parse JSON
                result = json.loads(content)
                elapsed = time.time() - start_time

                st.success("✅ Fertilization Plan Generated Successfully!")
                st.write(f"**Processing Time:** {elapsed:.2f} seconds")

                # Display result nicely
                st.subheader("📦 Fertilization Plan JSON Output")
                st.json(result)

                # Download button
                result_text = json.dumps(result, indent=2, ensure_ascii=False)
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"fertilization_plan_{timestamp}.txt"

                st.download_button(
                    label="💾 Download Result (.txt)",
                    data=result_text.encode("utf-8"),
                    file_name=filename,
                    mime="text/plain"
                )

                # Copyable output section
                st.subheader("📋 Copyable JSON Output")
                st.code(result_text, language="json")

            except json.JSONDecodeError as e:
                st.error(f"⚠️ JSON Parsing Error: {e}")
                st.text_area("Raw Response", content, height=300)
            except Exception as e:
                st.error(f"❌ Error: {e}")
                st.exception(e)