import os, io, json, time, pickle, hashlib
from typing import List, Dict, Tuple
import gradio as gr
import numpy as np

# ============ Lightweight local model (CPU) ============
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

LLM_ID = os.environ.get("LLM_ID", "google/flan-t5-base")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "256"))

_tokenizer = AutoTokenizer.from_pretrained(LLM_ID)
_model = AutoModelForSeq2SeqLM.from_pretrained(LLM_ID)

def llm(prompt: str, temperature: float = 0.2):
    inputs = _tokenizer(prompt, return_tensors="pt")
    outputs = _model.generate(
        **inputs,
        do_sample=temperature > 0.0,
        temperature=temperature,
        max_new_tokens=MAX_NEW_TOKENS,
        num_beams=1 if temperature > 0 else 4,
    )
    return _tokenizer.decode(outputs[0], skip_special_tokens=True)

# ============ Tiny RAG store (local, CPU) ============
from sentence_transformers import SentenceTransformer

EMBED_ID = os.environ.get("EMBED_ID", "sentence-transformers/all-MiniLM-L6-v2")
_embedder = SentenceTransformer(EMBED_ID)

DB_DIR = os.environ.get("DB_DIR", "./data/agri_space_db")
os.makedirs(DB_DIR, exist_ok=True)

def _db_path(ns: str) -> str:
    h = hashlib.sha1(ns.encode()).hexdigest()[:10]
    return os.path.join(DB_DIR, f"rag_{h}.pkl")

def _load_db(ns: str):
    p = _db_path(ns)
    if os.path.exists(p):
        with open(p, "rb") as f:
            return pickle.load(f)
    return {"chunks": [], "embeds": None}

def _save_db(ns: str, db: Dict):
    with open(_db_path(ns), "wb") as f:
        pickle.dump(db, f)

def chunk_text(text: str, max_tokens:int=220, overlap:int=40) -> List[str]:
    words = text.split()
    chunks = []
    i = 0
    step = max_tokens - overlap
    while i < len(words):
        chunk = " ".join(words[i:i+max_tokens])
        chunks.append(chunk)
        i += step
    return chunks

def add_docs(namespace: str, docs: List[Tuple[str, str]]) -> str:
    db = _load_db(namespace)
    new_chunks = []
    for name, text in docs:
        for ch in chunk_text(text):
            new_chunks.append({"name": name, "text": ch})
    if not new_chunks:
        return "No text detected."
    embeds = _embedder.encode([c["text"] for c in new_chunks], normalize_embeddings=True)
    if db["chunks"]:
        db["chunks"].extend(new_chunks)
        if db["embeds"] is not None:
            db["embeds"] = np.vstack([db["embeds"], embeds])
        else:
            db["embeds"] = embeds
    else:
        db["chunks"] = new_chunks
        db["embeds"] = embeds
    _save_db(namespace, db)
    return f"Added {len(new_chunks)} chunks to namespace '{namespace}'."

def retrieve(namespace: str, question: str, top_k:int=5) -> List[Dict]:
    db = _load_db(namespace)
    if not db["chunks"]:
        return []
    qv = _embedder.encode([question], normalize_embeddings=True)[0]
    sims = db["embeds"] @ qv
    idx = np.argsort(-sims)[:top_k]
    results = []
    for i in idx:
        ch = db["chunks"][int(i)]
        results.append({"name": ch["name"], "text": ch["text"], "score": float(sims[int(i)])})
    return results

def rag_answer(namespace: str, question: str, temperature: float=0.2, top_k:int=5) -> Tuple[str, str]:
    ctx = retrieve(namespace, question, top_k=top_k)
    if not ctx:
        return ("No documents indexed yet. Add farm PDFs/text first.", "")
    context_text = "\n\n".join([f"[{i+1}] {c['name']}: {c['text'][:600]}" for i,c in enumerate(ctx)])
    prompt = f"""You are AgriTech Copilot for vertical agriculture, homesteading, off-grid power, regenerative farming, and preparedness.
Use ONLY the context when answering. Be practical and safe. Cite the snippets you used by bracket numbers.

Question: {question}

Context:
{context_text}

Answer with steps, materials, and cautions where relevant. Include citations like [1], [2]."""
    ans = llm(prompt, temperature=temperature)
    sources = "\n".join([f"[{i+1}] {c['name']} (score {c['score']:.3f})" for i,c in enumerate(ctx)])
    return ans, sources

# ============ Agri/Off-grid calculators ============
def seed_bed_planner(crop:str, bed_length_ft:float, bed_width_ft:float, in_row_spacing_in:float, row_spacing_in:float):
    bed_len_in = bed_length_ft*12.0
    bed_w_in  = bed_width_ft*12.0
    rows = max(1, int(bed_w_in // row_spacing_in))
    plants_per_row = max(1, int(bed_len_in // in_row_spacing_in))
    total_plants = rows * plants_per_row
    return {
        "rows": rows,
        "plants_per_row": plants_per_row,
        "total_plants": total_plants,
        "note": f"Estimate for {crop}. Adjust for cultivar and local climate."
    }

def solar_offgrid_sizer(daily_wh:float, sun_hours:float=5.0, system_voltage:float=24.0):
    panel_watts = daily_wh / sun_hours * 1.25
    battery_wh = daily_wh * 2.0
    battery_ah = battery_wh / system_voltage
    inverter_w = daily_wh / 24 * 3
    return {
        "suggested_pv_watts": round(panel_watts, 1),
        "battery_wh": round(battery_wh, 1),
        "battery_ah_at_voltage": round(battery_ah, 1),
        "inverter_watts_peak": round(inverter_w, 1),
        "note": "Sizing is conservative. Validate with a local installer. Use fuses, disconnects, and proper wire gauge."
    }

def survival_checklist(days:int=7, people:int=2):
    water_gal = people * days * 1.0
    calories = people * days * 2000
    items = [
        f"Water: ~{water_gal} gallons",
        f"Food: ~{calories} kcal",
        "Heat/cook: camp stove + fuel",
        "Light: headlamps + batteries",
        "Sanitation: bags, bleach, soap",
        "Med kit: prescriptions, dressings, tape",
        "Tools: multitool, cordage, tarp",
        "Comms: power bank, radio",
        "Docs: IDs, cash, copies"
    ]
    return "\n".join(items)

# ============ Gradio UI ============
with gr.Blocks(title="Vertical AI: Agriculture - Homesteading - Off-Grid - Preparedness") as demo:
    gr.Markdown("## Vertical AI Copilot — Agriculture · Homesteading · Off-Grid · Preparedness")

    # ---------------- Agri Copilot (Chat) ----------------
    with gr.Tab("Agri Copilot (Chat)"):
        chat_in = gr.Textbox(label="Question", lines=3)
        temp = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Creativity")
        out = gr.Textbox(label="Response", lines=12)
        ask_btn = gr.Button("Answer")

        def answer_plain(q, t):
            prompt = "You are a practical agriculture/off-grid assistant.\n\nQ: {}\nA:".format(q)
            return llm(prompt, temperature=t)

        ask_btn.click(answer_plain, [chat_in, temp], [out])

    # ---------------- Knowledge (RAG) ----------------
    with gr.Tab("Knowledge (RAG)"):
        ns2 = gr.Textbox(value="farm-commons", label="Namespace")
        up = gr.Files(label="Upload .txt/.md", file_types=["text"])
        add_btn = gr.Button("Add to Knowledge")
        add_out = gr.Textbox(label="Indexer Output")

        def add_action(namespace, files):
            docs = []
            if files:
                for f in files:
                    try:
                        # Name
                        name = os.path.basename(getattr(f, "name", "uploaded.txt"))

                        # Read as text safely
                        if hasattr(f, "read"):
                            content = f.read()
                            if isinstance(content, bytes):
                                content = content.decode("utf-8", errors="ignore")
                        else:
                            # f may be a path-like
                            with open(f, "r", encoding="utf-8", errors="ignore") as fh:
                                content = fh.read()

                        docs.append((name, content))
                    except Exception as e:
                        return f"Error reading {name}: {e}"

            if not docs:
                return "No files processed."
            return add_docs(namespace, docs)

        add_btn.click(add_action, [ns2, up], [add_out])

        q2 = gr.Textbox(label="Question", lines=3)
        k = gr.Slider(1, 10, value=5, step=1, label="Top-K")
        t2 = gr.Slider(0.0, 1.0, value=0.2, step=0.05, label="Creativity")
        ask2 = gr.Button("Answer with RAG")
        ans2 = gr.Textbox(label="Answer", lines=12)
        src2 = gr.Textbox(label="Sources", lines=8)

        ask2.click(
            lambda ns, q, t, topk: rag_answer(ns, q, t, int(topk)),
            [ns2, q2, t2, k],
            [ans2, src2]
        )

    # ---------------- Tools ----------------
    with gr.Tab("Tools"):
        # Seed/Bed Planner
        crop = gr.Textbox(label="Crop", value="Carrot")
        bl = gr.Number(label="Bed length (ft)", value=50)
        bw = gr.Number(label="Bed width (ft)", value=2.5)
        ir = gr.Number(label="In-row spacing (in)", value=2.0)
        rs = gr.Number(label="Row spacing (in)", value=6.0)
        calc1 = gr.Button("Calculate")
        out1 = gr.JSON(label="Plan")
        calc1.click(seed_bed_planner, [crop, bl, bw, ir, rs], [out1])

        # Off-Grid Solar Sizer
        dwh = gr.Number(label="Daily energy use (Wh/day)", value=3000)
        sun = gr.Number(label="Sun hours/day", value=5.0)
        volt = gr.Number(label="System voltage (V)", value=24.0)
        calc2 = gr.Button("Size System")
        out2 = gr.JSON(label="Sizing")
        calc2.click(solar_offgrid_sizer, [dwh, sun, volt], [out2])

        # Survival Checklist
        days = gr.Number(label="Days", value=7)
        ppl = gr.Number(label="People", value=2)
        calc3 = gr.Button("Build Checklist")
        out3 = gr.Textbox(label="Checklist", lines=10)
        calc3.click(survival_checklist, [days, ppl], [out3])

    # ---------------- About ----------------
    with gr.Tab("About & Safety"):
        gr.Markdown(
            "**Safety:** General guidance only. Follow local laws, building/electrical codes, and best practices."
        )

demo.launch()