import google.generativeai as genai_generative
import json
import time

# Configure API key
genai_generative.configure(api_key="AIzaSyA9jCSJ6HTwoDeEDb4mO7dmh15dhQ64Kqc")

# Defect classification mapping
issue_map = """

Defect Description: > Hairline multiple shallow cracks.Small multiple cracks that looks like spider web.
Defect Class: RCC - Shrinkage cracks

Defect Description: > Found on load-bearing members like column, beam, slab Have more depth then hairline crack Have width of more than 0.3 mm.
Defect Class: RCC - Crack

Defect Description: > Detaching layer of concrete along the reinforcement / TMT. Exposed rusted reinforcement / TMT. Hollowness on concrete surface.
Defect Class: RCC - Spalling of concrete

Defect Description: > Hairline multiple shallow cracks. Small multiple cracks that looks like spider web.
Defect Class: Plaster - Shrinkage Cracks

Defect Description: > Found on RCC & brick / block wall junction Generally a straight crack (Vertical or horizontal)
Defect Class: Plaster - Separation Cracks

Defect Description: > Generally found on the middle of the wall. Either stepped or diagonal is nature.  
Defect Class: Plaster - Diagonal crack with less than 3 mm width

Defect Description: > Generally found on the middle of the wall. Either stepped or diagonal is nature.  
Defect Class: Plaster - Diagonal crack with more than 3 mm width

Defect Description: >Looks like uneven & patchy at places. Uneven gap noted between plaster & straight edge.
Defect Class: Plaster - Uneven Surface

Defect Description: > Loose sand comes out of plaster when rubbed with solid object.
Defect Class: Plaster - Loose sand

Defect Description: > Appear like loose white powder over plaster surface. Generally appear on moist / wet surface. Common with bed bricks walls.
Defect Class: Plaster - Efflorescence

Defect Description: > Hollow sound comes out of plaster when tapped with metal object / hollow stick. Delaminating of plaster layer from brick or RCC surface.
Defect Class: Plaster - Deboning (Hollowness)

"""

# Prompt builders
def _build_prompt(image_description: str = "") -> str:
    if image_description:
        return (
            "Based on the image and image description, match the defect description, "
            "and classify the defect into the correct defect class.\n\n"
            f"Image Description: {image_description}\n\n"
            "Defect Description and class: \n"
            f"{issue_map}\n\n"
            "format the response in the following format:\n"
            "Defect Class: <defect_class>\n"
            "Reasoning: <reasoning>"
        )
    return (
        " Based on the image and the defect description, \n"
        "please classify the defect into the correct defect class.\n"
        "Defect Description and class: \n\n"
        f"{issue_map}\n\n"
        "format the response in the following format:\n"
        "Defect Class: <defect_class>\n"
        "Reasoning: <reasoning>"
    )

# Initialize model
model = genai_generative.GenerativeModel("gemini-2.0-flash", generation_config={"temperature": 0.5})

def upload_file(image_path, image_description: str = ""):
    """Upload image to Gemini and get classification response"""
    myfile = genai_generative.upload_file(image_path)
    while myfile.state.name == "PROCESSING":
        time.sleep(0.5)
        myfile = genai_generative.get_file(myfile.name)

    try:
        prompt = _build_prompt(image_description)
        image_response = model.generate_content(contents=[prompt, myfile])
        print(image_response.text)
    finally:
        # Clean up the uploaded file to prevent gRPC timeout warnings
        try:
            genai_generative.delete_file(myfile.name)
        except Exception as e:
            print(f"Warning: Could not delete uploaded file: {e}")

    return image_response.text

# Load GKA.json knowledge base
with open('GKA.json', 'r') as file:
    gka_data = json.load(file)

gka_data_str = json.dumps(gka_data, indent=2)

def chatbot(image_path, image_description: str = ""):
    """Simple chatbot with max 2 interactions: image analysis + 1 follow-up question"""
    
    print("=== Building Defect Analysis Chatbot ===")
    print("Step 1: Analyzing image for defects...")
    
    # Step 1: Analyze image
    image_response = upload_file(image_path, image_description)
    
    print("\n=== Image Analysis Complete ===")
    print("Defect identified! You can ask 1 follow-up question.")
    
    # Step 2: One follow-up question
    query = input("\nEnter your question about this defect: ")
    
    if query.strip():
        print("\n=== Answering your question ===")
        
        # Extract defect class and reasoning from image response
        try:
            defect_class = image_response.split("Defect Class: ")[1].split("\n")[0] if "Defect Class: " in image_response else "Unknown"
            reasoning = image_response.split("Reasoning: ")[1] if "Reasoning: " in image_response else "No reasoning provided"
        except:
            defect_class = "Unknown"
            reasoning = "No reasoning provided"
        
        # Generate answer based on knowledge base
        prompt = f"""Based on the defect analysis, answer this question using the knowledge base.

Defect Class: {defect_class}
Reasoning: {reasoning}
Question: {query}

Knowledge Base: {gka_data_str}

Provide a helpful answer based on the knowledge base information."""
        
        answer = model.generate_content(contents=[prompt])
        print(f"\nAnswer: {answer.text}")
        
        return {
            "defect_class": defect_class,
            "reasoning": reasoning,
            "question": query,
            "answer": answer.text
        }
    else:
        print("No question asked. Chatbot session complete.")
        return {
            "defect_class": image_response.split("Defect Class: ")[1].split("\n")[0] if "Defect Class: " in image_response else "Unknown",
            "reasoning": image_response.split("Reasoning: ")[1] if "Reasoning: " in image_response else "No reasoning provided",
            "question": None,
            "answer": None
        }

if __name__ == "__main__":
    # Simple chatbot usage
    image_path = "1000039877.jpeg"
    result = chatbot(image_path)
    
    print("\n=== Chatbot Session Summary ===")
    print(f"Defect Class: {result['defect_class']}")
    print(f"Reasoning: {result['reasoning']}")
    if result['question']:
        print(f"Question: {result['question']}")
        print(f"Answer: {result['answer']}")
    else:
        print("No follow-up question was asked.")