import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from sentence_transformers import SentenceTransformer
import joblib
import pandas as pd
from huggingface_hub import hf_hub_download
import joblib

# ----------------------------
# LOAD MODELS
# ----------------------------

schema_model_id = "gouravkumar23/schema-recommendation-model"
evolution_model_id = "gouravkumar23/schema-evolution-predictor"
doc_model_id = "gouravkumar23/schema-documentation-generator"

schema_tokenizer = T5Tokenizer.from_pretrained(schema_model_id)
schema_model = T5ForConditionalGeneration.from_pretrained(schema_model_id)

evo_tokenizer = GPT2Tokenizer.from_pretrained(evolution_model_id)
evo_model = GPT2LMHeadModel.from_pretrained(evolution_model_id)

doc_tokenizer = T5Tokenizer.from_pretrained(doc_model_id)
doc_model = T5ForConditionalGeneration.from_pretrained(doc_model_id)

semantic_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

model_path = hf_hub_download(
    repo_id="gouravkumar23/column-type-inference",
    filename="model.pkl"
)

type_model = joblib.load(model_path)

# ----------------------------
# FUNCTIONS
# ----------------------------

def recommend_schema(schema):

    prompt = "schema history: " + schema

    inputs = schema_tokenizer(prompt, return_tensors="pt")

    outputs = schema_model.generate(**inputs, max_length=40)

    return schema_tokenizer.decode(outputs[0], skip_special_tokens=True)


def predict_schema(schema):

    prompt = schema + " ->"

    inputs = evo_tokenizer(prompt, return_tensors="pt")

    outputs = evo_model.generate(**inputs, max_length=40)

    return evo_tokenizer.decode(outputs[0], skip_special_tokens=True)


def generate_doc(column):

    prompt = "describe column: " + column

    inputs = doc_tokenizer(prompt, return_tensors="pt")

    outputs = doc_model.generate(**inputs, max_length=40)

    return doc_tokenizer.decode(outputs[0], skip_special_tokens=True)


def infer_type(value):

    value = str(value)

    features = {
        "length": len(value),
        "is_digit": value.isdigit(),
        "has_dash": "-" in value,
        "has_decimal": "." in value,
        "is_alpha": value.isalpha()
    }

    df = pd.DataFrame([features])

    return type_model.predict(df)[0]

# ----------------------------
# GRADIO INTERFACE
# ----------------------------

schema_ui = gr.Interface(
    fn=recommend_schema,
    inputs="text",
    outputs="text",
    title="Schema Recommendation Engine"
)

evolution_ui = gr.Interface(
    fn=predict_schema,
    inputs="text",
    outputs="text",
    title="Schema Evolution Predictor"
)

doc_ui = gr.Interface(
    fn=generate_doc,
    inputs="text",
    outputs="text",
    title="Column Documentation Generator"
)

type_ui = gr.Interface(
    fn=infer_type,
    inputs="text",
    outputs="text",
    title="Column Type Inference"
)

app = gr.TabbedInterface(
    [schema_ui, evolution_ui, doc_ui, type_ui],
    ["Recommend", "Predict", "Docs", "Type"]
)

app.launch()