app.py

"""
Farm Drip Irrigation Design Tool - Phase 1.5 (Geometry + Valve Engine)

Two modes:
1. Quick test: Type geofence coordinates, select crop, get design
2. Full pipeline: Upload GeoJSON from user app, run valve + drip layout, download result

Supports:
- GeoJSON input/output (RFC 7946 standard)
- Valve placement with hierarchical decision matrix
- Drip layout generation per valve zone
- BOM estimation
"""

import os
os.environ["GRADIO_SSR_MODE"] = "false"  # Disable Node.js SSR so custom FastAPI routes work

import gradio as gr
import json
import math
from shapely.geometry import Polygon
from PIL import Image, ImageDraw
from drip_engine import (
    parse_geofence_to_polygon,
    validate_polygon,
    generate_drip_layout,
    estimate_bom,
    design_summary,
    DripLayoutError,
)
from design_api import process_farm_design
from unit_converter import supported_length_units, supported_area_units, length_to_meters
from llm_chat import FarmerAssistant
from farm_data_collector import FarmQACollector, FarmData

DEFAULT_GEOFENCE = "0,0;100,0;100,100;0,100"
DEFAULT_CROP = "generic"
DEFAULT_HEADLAND = 1.0

# Initialize LLM assistant (lazy-loaded to avoid blocking if HF_TOKEN not set)
_ASSISTANT = None

def get_assistant() -> FarmerAssistant:
    """Lazy-load the assistant on first use."""
    global _ASSISTANT
    if _ASSISTANT is None:
        try:
            _ASSISTANT = FarmerAssistant()
        except ValueError as e:
            # Let the UI handle the error
            raise e
    return _ASSISTANT

# Session storage for farm data collection per chat
_COLLECTORS = {}  # Map from session_id to FarmQACollector

def get_or_create_collector(session_id: str) -> FarmQACollector:
    """Get or create a collector for this chat session."""
    if session_id not in _COLLECTORS:
        _COLLECTORS[session_id] = FarmQACollector()
    return _COLLECTORS[session_id]

# Build dropdown choices from the registry — label (canonical key)
_COORD_UNIT_CHOICES = [
    f"{label} ({key})" for key, label in supported_length_units().items()
]
_AREA_UNIT_CHOICES = [
    f"{label} ({key})" for key, label in supported_area_units().items()
]

def _parse_unit_choice(choice: str) -> str:
    """Extract canonical key from a 'Label (key)' dropdown string."""
    return choice.split("(")[-1].rstrip(")")


def generate_design(geofence_text, crop, headland_m, override_spacing,
                    coord_unit_choice, area_unit_choice):
    """Quick test mode with text geofence."""
    try:
        coord_unit = _parse_unit_choice(coord_unit_choice)
        area_unit  = _parse_unit_choice(area_unit_choice)

        polygon_px = parse_geofence_to_polygon(geofence_text, coord_unit=coord_unit)
        is_valid, msg = validate_polygon(polygon_px)
        if not is_valid:
            return None, "{}", msg, f"❌ Invalid polygon: {msg}"

        override_sp = override_spacing if override_spacing > 0 else None
        design = generate_drip_layout(
            polygon_px,
            crop=crop,
            headland_buffer_m=headland_m,
            override_spacing_m=override_sp,
        )
        bom = estimate_bom(design, unit="usd")
        summary = design_summary(design, bom, area_unit=area_unit)
        image = _render_design_image(design, polygon_px)
        bom_json = json.dumps(bom, indent=2)
        return image, bom_json, summary, "✅ Design generated successfully"

    except DripLayoutError as e:
        return None, "{}", "", f"❌ Design error: {str(e)}"
    except Exception as e:
        return None, "{}", "", f"❌ Unexpected error: {str(e)}"


def _render_design_image(design, polygon_px):
    """Render the drip layout on a canvas image."""
    minx, miny, maxx, maxy = polygon_px.bounds
    width_px = max(int(maxx - minx) + 20, 400)
    height_px = max(int(maxy - miny) + 20, 300)

    image = Image.new("RGB", (width_px, height_px), color="white")
    draw = ImageDraw.Draw(image)
    scale = min((width_px - 40) / (maxx - minx), (height_px - 40) / (maxy - miny))
    offset_x = 20 - minx * scale
    offset_y = 20 - miny * scale

    def scale_point(x, y):
        return (x * scale + offset_x, y * scale + offset_y)

    boundary_coords = list(polygon_px.exterior.coords)
    scaled_boundary = [scale_point(x, y) for x, y in boundary_coords]
    if len(scaled_boundary) > 2:
        draw.polygon(scaled_boundary, outline="green", width=3)

    main_line = design["main_line"]
    main_coords = [scale_point(x, y) for x, y in main_line.coords]
    if len(main_coords) > 1:
        draw.line(main_coords, fill="red", width=4)

    for lateral in design["laterals"]:
        lateral_coords = [scale_point(x, y) for x, y in lateral.coords]
        if len(lateral_coords) > 1:
            draw.line(lateral_coords, fill="blue", width=2)

    draw.text((10, 10), f"Farm Drip Layout - {design['crop'].title()}", fill="black")
    draw.text((10, 30), f"Area: {design['farm_area_ha']:.2f} ha", fill="black")
    draw.text((10, 50), f"Emitters: {design['emitter_count']}", fill="black")

    return image


def process_geojson_input(geojson_text):
    """Full pipeline: accept GeoJSON string, return output + summary."""
    try:
        result = process_farm_design(geojson_text)

        props = result.get("properties", {})
        if props.get("type") == "farm_design_error":
            err = props.get("error", {})
            return json.dumps(result, indent=2), "", None, f"❌ {err.get('code')}: {err.get('message')}"

        summary_lines = []
        ds = props.get("design_summary", {})
        summary_lines.append("=== Farm Design Summary ===")
        summary_lines.append(f"Farm Area:       {ds.get('farm_area_ha', 'N/A')} ha")
        summary_lines.append(f"Total Valves:    {ds.get('total_valves', 'N/A')}")
        summary_lines.append(f"Drip Tape:       {ds.get('total_drip_tape_m', 'N/A')} m")
        summary_lines.append(f"Main Line:       {ds.get('total_main_line_m', 'N/A')} m")
        summary_lines.append(f"Emitters:        {ds.get('total_emitters', 'N/A')}")
        summary_lines.append(f"Pump:            {ds.get('pump_hp', 'N/A')} HP ({ds.get('pump_flow_lph', 'N/A')} L/h)")
        summary_lines.append(f"Strategy:        {ds.get('manifold_strategy', 'N/A')}")
        summary_lines.append("")

        bom = props.get("bom", {})
        summary_lines.append("=== Bill of Materials ===")
        summary_lines.append(f"Main Pipe (16mm):  {bom.get('main_line_16mm_m', 'N/A')} m")
        summary_lines.append(f"Drip Tape (16mm):  {bom.get('drip_tape_16mm_m', 'N/A')} m")
        summary_lines.append(f"Inline Emitters:   {bom.get('inline_emitters', 'N/A')}")
        summary_lines.append(f"Valves:            {bom.get('valves_count', 'N/A')}")
        if "total_cost_usd" in bom:
            summary_lines.append(f"Total Cost:        ${bom.get('total_cost_usd', 'N/A')}")
        summary_lines.append("")

        zones = props.get("zone_details", [])
        if zones:
            summary_lines.append("=== Per-Zone Breakdown ===")
            for z in zones:
                if "error" in z:
                    summary_lines.append(f"  {z['valve_id']}: {z['crop']} — ERROR: {z['error']}")
                else:
                    summary_lines.append(
                        f"  {z['valve_id']}: {z['crop']} | "
                        f"{z.get('area_ha', 0):.2f} ha | "
                        f"{z.get('emitters', 0)} emitters | "
                        f"{z.get('lateral_m', 0):.0f} m tape"
                    )

        summary = "\n".join(summary_lines)

        feature_types = {}
        for f in result.get("features", []):
            t = f.get("properties", {}).get("type", "unknown")
            feature_types[t] = feature_types.get(t, 0) + 1
        visual_text = "Features generated:\n" + "\n".join(f"  {k}: {v}" for k, v in feature_types.items())

        output_json = json.dumps(result, indent=2)
        return output_json, summary, visual_text, "✅ Design generated from GeoJSON"

    except Exception as e:
        return "{}", "", "", f"❌ Error: {str(e)}"


def process_geojson_file(file_obj):
    """Accept uploaded JSON file, run pipeline."""
    if file_obj is None:
        return "{}", "", None, "Please upload a GeoJSON file."
    try:
        with open(file_obj.name, "r", encoding="utf-8") as f:
            content = f.read()
        return process_geojson_input(content)
    except Exception as e:
        return "{}", "", None, f"❌ File error: {str(e)}"





# ── Shared pipeline helpers ───────────────────────────────────────────

def _build_farm_geojson(coords_lonlat, pump_hp, crop, headland_m):
    """
    Wrap a list of [lon, lat] coordinates into a minimal GeoJSON
    FeatureCollection that process_farm_design() can consume.
    """
    ring = coords_lonlat + [coords_lonlat[0]]   # close the ring
    return json.dumps({
        "type": "FeatureCollection",
        "properties": {
            "pump_hp": float(pump_hp),
            "headland_buffer_m": float(headland_m),
        },
        "features": [
            {"type": "Feature",
             "properties": {"type": "farm_boundary"},
             "geometry": {"type": "Polygon", "coordinates": [ring]}},
            {"type": "Feature",
             "properties": {"type": "pump", "pump_hp": float(pump_hp)},
             "geometry": {"type": "Point", "coordinates": ring[0]}},
            {"type": "Feature",
             "properties": {"type": "crop_zone", "crop": crop},
             "geometry": {"type": "Polygon", "coordinates": [ring]}},
        ],
    })


def design_from_form(center_lat, center_lon, field_width, field_height,
                    dim_unit_choice, pump_hp, crop, headland_m):
    """
    Simple Form tab: build a rectangular farm from centre lat/lon +
    width/height, then run the full GeoJSON design pipeline.
    """
    try:
        dim_unit = _parse_unit_choice(dim_unit_choice)
        w_m = length_to_meters(float(field_width),  dim_unit)
        h_m = length_to_meters(float(field_height), dim_unit)
        lat, lon = float(center_lat), float(center_lon)

        # Flat-earth degree offsets (accurate enough for farm-scale areas)
        lat_deg = 1.0 / 111_320.0
        lon_deg = 1.0 / (111_320.0 * math.cos(math.radians(lat)))
        hw, hh  = (w_m / 2) * lon_deg, (h_m / 2) * lat_deg

        coords = [
            [lon - hw, lat - hh], [lon + hw, lat - hh],
            [lon + hw, lat + hh], [lon - hw, lat + hh],
        ]
        out_json, summary, _, status = process_geojson_input(
            _build_farm_geojson(coords, pump_hp, crop, headland_m)
        )
        return out_json, summary, status
    except Exception as e:
        return "{}", "", f"❌ Error: {e}"


def chat_with_assistant(user_message: str, chat_history: list) -> tuple:
    """
    Smart chat handler that either:
    1. Collects farm design data via Q&A, OR
    2. Answers general irrigation questions
    
    Args:
        user_message: The farmer's question or response
        chat_history: List of dicts with {'role': 'user'/'assistant', 'content': '...'}
        
    Returns:
        (updated_chat_history, status_message)
    """
    if not user_message.strip():
        return chat_history, "Please enter a message."
    
    try:
        # Use a stable session ID (in a real app, use request.session_id)
        # For now, use hash of initial message as session marker
        session_id = "chat_session"  # TODO: replace with real session management
        collector = get_or_create_collector(session_id)
        
        # First message - show welcome if not started
        if not collector.conversation_started and len(chat_history or []) == 0:
            welcome = collector.get_initial_prompt()
            updated_history = (chat_history or []) + [
                {"role": "assistant", "content": welcome}
            ]
            return updated_history, "✅ Welcome message sent"
        
        # Process user input through the smart collector
        response, farm_data = collector.process_user_input(user_message)
        
        # Update chat history
        updated_history = (chat_history or []) + [
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": response}
        ]
        
        # If farm data collection is complete, generate design automatically
        if farm_data and farm_data.is_complete():
            try:
                # Generate design from collected data
                design_json, design_summary, status = design_from_form(
                    farm_data.center_lat,
                    farm_data.center_lon,
                    farm_data.field_width,
                    farm_data.field_height,
                    "Meters (m)",  # default unit
                    farm_data.pump_hp,
                    farm_data.crop,
                    farm_data.headland_buffer_m or 1.0
                )
                
                # Add design result to chat
                design_msg = (
                    f"✅ **Design Generated!**\n\n"
                    f"{design_summary}\n\n"
                    f"Your complete farm design is ready. "
                    f"Check the 'Simple Form' tab to view the full GeoJSON output."
                )
                updated_history = updated_history + [
                    {"role": "assistant", "content": design_msg}
                ]
                return updated_history, "✅ Design generated and added to chat"
                
            except Exception as design_error:
                error_msg = f"⚠️ Error generating design: {str(design_error)}"
                updated_history = updated_history + [
                    {"role": "assistant", "content": error_msg}
                ]
                return updated_history, f"Design generation error: {design_error}"
        
        return updated_history, "✅ Response sent"
        
    except Exception as e:
        import traceback
        return chat_history, f"❌ Error: {str(e)}\n{traceback.format_exc()[:200]}"






# Load sample files for display
with open("samples/input_example.json", "r") as f:
    SAMPLE_INPUT = f.read()
with open("samples/output_example.json", "r") as f:
    SAMPLE_OUTPUT = f.read()


with gr.Blocks(title="Farm Drip Irrigation Designer") as demo:
    gr.Markdown(
        """
# 🌾 Farm Drip Irrigation Designer
**Design drip irrigation layouts for any field shape — no API key required.**

| Tab | Best for |
|-----|----------|
| **Simple Form** | Rectangular farms — enter centre location + dimensions |
| **Draw on Map** | Any shape — draw your field boundary on OpenStreetMap |
| **Quick Test** | Developers — paste raw coordinates |
| **GeoJSON Pipeline** | Apps — upload/paste full GeoJSON |
"""
    )

    with gr.Tabs():

        # ── Tab 1: Simple Form ────────────────────────────────────────
        with gr.TabItem("Simple Form"):
            gr.Markdown(
                "Enter your farm’s **centre coordinates** and **field dimensions**. "
                "Tip: right-click any point in Google Maps or OpenStreetMap to copy lat/lon."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Farm Location")
                    form_lat = gr.Number(label="Centre Latitude",  value=12.9716,
                                        info="Decimal degrees, e.g. 12.9716")
                    form_lon = gr.Number(label="Centre Longitude", value=77.5946,
                                        info="Decimal degrees, e.g. 77.5946")
                    gr.Markdown("### Field Size")
                    with gr.Row():
                        form_width  = gr.Number(label="Width",  value=200, minimum=1)
                        form_height = gr.Number(label="Height", value=150, minimum=1)
                    form_dim_unit = gr.Dropdown(
                        label="Dimension Unit",
                        choices=_COORD_UNIT_CHOICES,
                        value="Meters (m)",
                    )
                    gr.Markdown("### Farm Parameters")
                    form_pump_hp  = gr.Slider(label="Pump HP",
                                             minimum=0.5, maximum=20.0, value=5.0, step=0.5)
                    form_crop     = gr.Dropdown(
                        label="Crop",
                        choices=["tomato","pepper","lettuce","cucumber","orchard","generic"],
                        value="generic",
                    )
                    form_headland = gr.Slider(label="Headland Buffer (m)",
                                             minimum=0.0, maximum=10.0, value=1.0, step=0.5)
                    form_btn = gr.Button("Generate Design", variant="primary")

                with gr.Column(scale=1):
                    gr.Markdown("### Results")
                    form_status  = gr.Textbox(label="Status", interactive=False)
                    form_summary = gr.Textbox(label="Design Summary", lines=20, interactive=False)

            with gr.Row():
                form_json = gr.Code(label="Output GeoJSON", language="json", lines=15)

            form_btn.click(
                fn=design_from_form,
                inputs=[form_lat, form_lon, form_width, form_height,
                        form_dim_unit, form_pump_hp, form_crop, form_headland],
                outputs=[form_json, form_summary, form_status],
            )

            gr.Markdown("### Quick Examples")
            gr.Examples(
                examples=[
                    [12.9716, 77.5946, 200, 150, "Meters (m)",   5.0, "tomato",  1.0],
                    [28.6139, 77.2090, 5.0, 3.0, "Acres (acres)",3.0, "lettuce", 1.0],
                    [30.7333, 76.7794, 10, 8,   "Chains (chain)",2.0, "generic", 0.5],
                ],
                inputs=[form_lat, form_lon, form_width, form_height,
                        form_dim_unit, form_pump_hp, form_crop, form_headland],
                fn=design_from_form,
                outputs=[form_json, form_summary, form_status],
                cache_examples=False,
            )

        # ── Tab 2: Quick Test (raw coordinates) ───────────────────────
        with gr.TabItem("Quick Test"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Farm Boundary Input")
                    geofence_input = gr.Textbox(
                        label="Geofence (x,y;x,y;x,y...)",
                        value=DEFAULT_GEOFENCE,
                        lines=3,
                        info="Comma-separated coords, semicolon-separated points",
                    )
                    crop_dropdown = gr.Dropdown(
                        label="Crop Type",
                        choices=["tomato", "pepper", "lettuce", "cucumber", "orchard", "generic"],
                        value=DEFAULT_CROP,
                    )
                    headland_slider = gr.Slider(
                        label="Headland Buffer (m)",
                        minimum=0.0, maximum=10.0, value=DEFAULT_HEADLAND, step=0.5,
                    )
                    spacing_slider = gr.Slider(
                        label="Override Lateral Spacing (m)",
                        minimum=0.0, maximum=5.0, value=0.0, step=0.1,
                        info="Set to 0 to use crop default",
                    )
                    coord_unit_dd = gr.Dropdown(
                        label="Coordinate Unit",
                        choices=_COORD_UNIT_CHOICES,
                        value="Meters (m)",
                        info="Unit of the x,y values in your geofence string",
                    )
                    area_unit_dd = gr.Dropdown(
                        label="Display Area In",
                        choices=_AREA_UNIT_CHOICES,
                        value="Hectares (ha)",
                        info="Unit used in the design summary output",
                    )
                    generate_btn = gr.Button("Generate Design", variant="primary")

                with gr.Column(scale=1):
                    gr.Markdown("### Design Visualization")
                    design_image = gr.Image(label="Drip Layout", type="pil")
                    status_text = gr.Textbox(label="Status", interactive=False)

            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Bill of Materials (USD)")
                    bom_json = gr.Code(label="BOM (JSON)", language="json")
                with gr.Column():
                    gr.Markdown("### Design Summary")
                    summary_text = gr.Textbox(label="Summary", lines=15, interactive=False)

            generate_btn.click(
                fn=generate_design,
                inputs=[geofence_input, crop_dropdown, headland_slider, spacing_slider,
                        coord_unit_dd, area_unit_dd],
                outputs=[design_image, bom_json, summary_text, status_text],
            )

            gr.Markdown("### Example Geofences (copy & paste)")
            gr.Examples(
                examples=[
                    ["0,0;100,0;100,100;0,100",            "tomato", 1.0, 0.0, "Meters (m)",  "Hectares (ha)"],
                    ["0,0;656,0;656,164;0,164",            "lettuce",1.0, 0.0, "Feet (ft)",   "Acres (acres)"],
                    ["0,0;100,0;100,100;50,100;50,50;0,50","pepper", 1.5, 0.0, "Meters (m)",  "Acres (acres)"],
                ],
                inputs=[geofence_input, crop_dropdown, headland_slider, spacing_slider,
                        coord_unit_dd, area_unit_dd],
                fn=generate_design,
                outputs=[design_image, bom_json, summary_text, status_text],
                cache_examples=False,
            )

        with gr.TabItem("GeoJSON Pipeline"):
            gr.Markdown(
                """
### Upload farm data as GeoJSON

Upload a GeoJSON FeatureCollection with your farm boundary, pump location, and crop zones.
The engine will place valves, generate drip layouts per zone, and return a complete design.

**Required**: `farm_boundary` (Polygon), `pump` (Point with pump_hp), crop zones (Polygons)
**Output**: GeoJSON with valves, valve_zones, main_lines, laterals + BOM
"""
            )

            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Input")
                    geojson_file = gr.File(
                        label="Upload GeoJSON (.json)",
                        file_types=[".json", ".geojson"],
                    )
                    geojson_text = gr.Textbox(
                        label="Or paste GeoJSON here",
                        lines=10,
                        placeholder='{"type": "FeatureCollection", "features": [...]}',
                    )
                    with gr.Row():
                        run_file_btn = gr.Button("Run from File", variant="primary")
                        run_text_btn = gr.Button("Run from Text")

                with gr.Column(scale=1):
                    gr.Markdown("### Output")
                    output_status = gr.Textbox(label="Status", interactive=False)
                    output_visual = gr.Textbox(label="Features Generated", lines=6, interactive=False)

            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Output GeoJSON")
                    output_geojson = gr.Code(label="Result (GeoJSON)", language="json", lines=20)
                with gr.Column():
                    gr.Markdown("### Design Summary")
                    output_summary = gr.Textbox(label="Summary", lines=20, interactive=False)

            run_file_btn.click(
                fn=process_geojson_file,
                inputs=[geojson_file],
                outputs=[output_geojson, output_summary, output_visual, output_status],
            )

            run_text_btn.click(
                fn=process_geojson_input,
                inputs=[geojson_text],
                outputs=[output_geojson, output_summary, output_visual, output_status],
            )

            gr.Markdown("### Example Input")
            gr.Code(SAMPLE_INPUT, language="json", label="Sample Input")

            gr.Markdown("### Example Output")
            gr.Code(SAMPLE_OUTPUT, language="json", label="Sample Output")

        # ── Tab 4: Farmer Assistant (LLM Chat) ─────────────────────────
        with gr.TabItem("🤖 Farmer Assistant"):
            gr.Markdown(
                """
### PhyFarm Assistant — Ask About Your Farm Design

Got questions about your drip irrigation system? Your crops? Cost estimation?
The AI assistant is here to help with practical, actionable guidance.

**What you can ask:**
- "How often should I water tomatoes with this setup?"
- "What's the flow rate per emitter?"
- "Can I use this design for peppers instead?"
- "How do I troubleshoot low pressure?"
- "What's included in the bill of materials?"
"""
            )

            with gr.Row():
                with gr.Column(scale=1):
                    chat_history = gr.Chatbot(
                        label="Conversation",
                        scale=1,
                        height=400,
                    )
                    chat_input = gr.Textbox(
                        label="Your question",
                        placeholder="Type your question here...",
                        lines=2,
                    )
                    with gr.Row():
                        chat_submit = gr.Button("Send", variant="primary")
                        chat_clear = gr.Button("Clear History")
                    chat_status = gr.Textbox(label="Status", interactive=False)

            # Connect chat submission
            chat_submit.click(
                fn=chat_with_assistant,
                inputs=[chat_input, chat_history],
                outputs=[chat_history, chat_status],
            ).then(
                lambda: "",  # Clear the input box after sending
                outputs=[chat_input],
            )

            # Clear history button
            chat_clear.click(
                fn=lambda: ([], "History cleared"),
                outputs=[chat_history, chat_status],
            )

            gr.Markdown(
                "**Note:** Connect your HuggingFace API token by setting the `HF_TOKEN` environment variable."
                "\n[Get a free token here](https://huggingface.co/settings/tokens)"
            )


# Expose REST endpoints alongside the Gradio UI.
#
# With ssr_mode=False, Gradio serves the SPA and static assets from
# Python (no Node.js SSR layer).  This lets us use mount_gradio_app
# so custom FastAPI routes are handled directly — the SSR layer no
# longer intercepts them with a 405.
from fastapi import FastAPI  # noqa: E402
from rest_api import build_router  # noqa: E402

api = FastAPI(title="Farm Layout Model API")
api.include_router(build_router())
app = gr.mount_gradio_app(api, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)