import gradio as gr
import duckdb
import pandas as pd
from datetime import datetime
import random

# --- MCP Registry ---
class MCP:
    def __init__(self):
        self.agents = {}
        self.services = {}  # To map service types to agent names

    def register_agent(self, name, agent, services=None):
        self.agents[name] = agent
        if services:
            for service in services:
                if service not in self.services:
                    self.services[service] = []
                self.services[service].append(name)

    def get_agent(self, name):
        return self.agents.get(name)

    def find_agents_by_service(self, service_type):
        return self.services.get(service_type, [])

mcp_registry = MCP()

# --- Data Management Agent ---
class DataManagementAgent:
    def __init__(self, db_path="market_ai_duckdb.db"):
        self.conn = duckdb.connect(db_path)
        self._create_tables()

    def _create_tables(self):
        try:
            self.conn.execute("""
            CREATE TABLE IF NOT EXISTS demand_forecasts (
                item_name VARCHAR,
                location VARCHAR,
                forecast_date VARCHAR,
                predicted_demand INTEGER
            );
            """)
        except Exception as e:
            print(f"Error creating demand_forecasts table: {e}")

        try:
            self.conn.execute("""
            CREATE TABLE IF NOT EXISTS consumer_requests (
                id INTEGER,
                consumer_name VARCHAR,
                item_name VARCHAR,
                requested_qty INTEGER,
                unit VARCHAR,
                required_date VARCHAR
            );
            """)
        except Exception as e:
            print(f"Error creating consumer_requests table: {e}")

        try:
            self.conn.execute("""
            CREATE TABLE IF NOT EXISTS farmer_stock (
                id INTEGER,
                farmer_name VARCHAR,
                item_name VARCHAR,
                available_qty INTEGER,
                unit VARCHAR,
                quality_grade VARCHAR,
                available_date VARCHAR
            );
            """)
        except Exception as e:
            print(f"Error creating farmer_stock table: {e}")
        self.conn.commit()


    def insert(self, table_name, data):
        cursor = self.conn.cursor()
        placeholders = ', '.join(['?'] * len(data))
        query = f"INSERT INTO {table_name} VALUES ({placeholders})"
        try:
            cursor.execute(query, data)
            self.conn.commit()
        except Exception as e:
            print(f"Error inserting into {table_name}: {e}")

    def fetch(self, query):
        try:
            df = self.conn.execute(query).df()
            return df
        except Exception as e:
            print(f"Error fetching data: {e}")
            return pd.DataFrame()

data_agent = DataManagementAgent()
mcp_registry.register_agent("data_agent", data_agent, services=["data_storage"])

# --- Demand Forecasting Agent (Simplified) ---
class DemandForecastingAgent:
    def __init__(self, data_agent):
        self.data_agent = data_agent

    def predict_demand(self, item_name, location, forecast_date):
        predicted_demand = random.randint(10, 100)
        self.data_agent.insert("demand_forecasts", (item_name, location, forecast_date, predicted_demand))
        return f"Predicted demand for {item_name} in {location} on {forecast_date}: {predicted_demand}"

demand_forecaster = DemandForecastingAgent(data_agent)
mcp_registry.register_agent("demand_forecaster", demand_forecaster, services=["demand_forecasting"])

# --- Vendor Agent (Illustrative) ---
class VendorAgent:
    def __init__(self, mcp_registry):
        self.mcp_registry = mcp_registry

    def get_demand_forecast(self, item_name, location, forecast_date):
        forecasting_agents = self.mcp_registry.find_agents_by_service("demand_forecasting")
        if forecasting_agents:
            forecaster_name = forecasting_agents[0]
            forecaster = self.mcp_registry.get_agent(forecaster_name)
            return forecaster.predict_demand(item_name, location, forecast_date)
        else:
            return "No demand forecasting agent available."

vendor_agent = VendorAgent(mcp_registry)
mcp_registry.register_agent("vendor_agent", vendor_agent)

# --- Consumer Agent ---
class ConsumerAgent:
    def __init__(self, data_agent):
        self.data_agent = data_agent

    def submit_request(self, consumer_name, item_name, requested_qty, unit, required_date):
        data = (int(datetime.now().timestamp()), consumer_name, item_name, requested_qty, unit, required_date)
        self.data_agent.insert("consumer_requests", data)
        return "Consumer request submitted successfully."

consumer_agent = ConsumerAgent(data_agent)
mcp_registry.register_agent("consumer_agent", consumer_agent, services=["consumer_request_submission"])

# --- Farmer Agent ---
class FarmerAgent:
    def __init__(self, data_agent):
        self.data_agent = data_agent

    def submit_stock(self, farmer_name, item_name, available_qty, unit, quality_grade, available_date):
        data = (int(datetime.now().timestamp()), farmer_name, item_name, available_qty, unit, quality_grade, available_date)
        self.data_agent.insert("farmer_stock", data)
        return "Farmer stock submitted successfully."

farmer_agent = FarmerAgent(data_agent)
mcp_registry.register_agent("farmer_agent", farmer_agent, services=["farmer_stock_submission"])

# --- Matching Agent ---
class MatchingAgent:
    def __init__(self, data_agent):
        self.data_agent = data_agent

    def view_matches(self):
        query = """
            SELECT c.consumer_name, c.item_name, c.requested_qty, c.unit AS consumer_unit, c.required_date,
                   f.farmer_name, f.available_qty, f.unit AS farmer_unit, f.quality_grade, f.available_date
            FROM consumer_requests c
            LEFT JOIN farmer_stock f ON c.item_name = f.item_name AND c.unit = f.unit
            WHERE c.required_date <= f.available_date
        """
        df = self.data_agent.fetch(query)
        return df

matching_agent = MatchingAgent(data_agent)
mcp_registry.register_agent("matching_agent", matching_agent, services=["market_matching"])

# --- Pre-Booking Agent ---
class PreBookingAgent:
    def pre_book_item(self, consumer_name, item_name):
        return f"{consumer_name} pre-booked {item_name}. Please proceed to checkout during market day."

prebooking_agent = PreBookingAgent()
mcp_registry.register_agent("prebooking_agent", prebooking_agent, services=["pre_booking"])

# --- GRADIO UI ---
def submit_consumer_ui(consumer_name, item_name, requested_qty, unit, required_date):
    agent = mcp_registry.get_agent("consumer_agent")
    return agent.submit_request(consumer_name, item_name, requested_qty, unit, required_date)

def submit_farmer_ui(farmer_name, item_name, available_qty, unit, quality_grade, available_date):
    agent = mcp_registry.get_agent("farmer_agent")
    return agent.submit_stock(farmer_name, item_name, available_qty, unit, quality_grade, available_date)

def view_market_matches_ui():
    agent = mcp_registry.get_agent("matching_agent")
    return agent.view_matches()

def pre_book_item_ui(consumer_name, item_name):
    agent = mcp_registry.get_agent("prebooking_agent")
    return agent.pre_book_item(consumer_name, item_name)

def get_demand_forecast_ui(item_name, location, forecast_date):
    vendor = mcp_registry.get_agent("vendor_agent")
    return vendor.get_demand_forecast(item_name, location, forecast_date)

with gr.Blocks() as app:
    gr.Markdown("## Agentic AI Market App with MCP")

    with gr.Tab("Consumer Request"):
        uid = gr.Textbox(label="Consumer Name")
        item = gr.Textbox(label="Item Name")
        qty = gr.Number(label="Quantity Required")
        unit = gr.Textbox(label="Unit (e.g., kg, liters)")
        req_date = gr.Textbox(label="Required Date (YYYY-MM-DD)")
        consumer_btn = gr.Button("Submit Request")
        consumer_output = gr.Textbox()
        consumer_btn.click(submit_consumer_ui, inputs=[uid, item, qty, unit, req_date], outputs=consumer_output)

    with gr.Tab("Farmer Stock"):
        fid = gr.Textbox(label="Farmer Name")
        item_stock = gr.Textbox(label="Item Name")
        qty_stock = gr.Number(label="Available Quantity")
        unit_stock = gr.Textbox(label="Unit (e.g., kg, liters)")
        quality = gr.Textbox(label="Quality Grade")
        avail_date = gr.Textbox(label="Available Date (YYYY-MM-DD)")
        farmer_btn = gr.Button("Submit Stock")
        farmer_output = gr.Textbox()
        farmer_btn.click(submit_farmer_ui, inputs=[fid, item_stock, qty_stock, unit_stock, quality, avail_date], outputs=farmer_output)

    with gr.Tab("Market Matches"):
        match_btn = gr.Button("View Matches")
        match_output = gr.Dataframe()
        match_btn.click(view_market_matches_ui, outputs=match_output)

    with gr.Tab("Pre-Booking"):
        consumer_name_prebook = gr.Textbox(label="Consumer Name")
        item_name_prebook = gr.Textbox(label="Item Name")
        prebook_btn = gr.Button("Pre-Book Item")
        prebook_output = gr.Textbox()
        prebook_btn.click(pre_book_item_ui, inputs=[consumer_name_prebook, item_name_prebook], outputs=prebook_output)

    with gr.Tab("Vendor Services"):
        item_forecast = gr.Textbox(label="Item Name")
        location_forecast = gr.Textbox(label="Location")
        date_forecast = gr.Textbox(label="Forecast Date (YYYY-MM-DD)")
        forecast_btn = gr.Button("Get Demand Forecast")
        forecast_output = gr.Textbox()
        forecast_btn.click(get_demand_forecast_ui, inputs=[item_forecast, location_forecast, date_forecast], outputs=forecast_output)

app.launch()