from copy import copy

from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.output_parsers import PydanticOutputParser
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from typing import Optional
import pandas as pd
import os

class InsuredDetails(BaseModel):
    """Insured business details"""

    business_name: Optional[str] = Field(description="The insured business name")
    business_address: str = Field(description="The insured business address")
    contact_name: Optional[str] = Field(description="the person of contact's name")
    contact_phone: Optional[str] = Field(description="the person of contact's phone number")

    def __str__(self):
        content = (f"business_name: {self.business_name}\nbusiness_address: {self.business_address}\n"
                   f"contact_name: {self.contact_name}\ncontact_phone: {self.contact_phone}")
        separator = "~" * 20
        return "\n".join([content, separator])


def parse_to_df(results):
    return pd.DataFrame([ins.model_dump() for ins in results])


load_dotenv()  # Only needed locally


google_api_key = os.getenv("GOOGLE_API_KEY")
if not google_api_key:
    raise ValueError("GOOGLE_API_KEY not set. Please check your environment variables.")

def llm_parser(docs, model="gemini-2.0-flash-001"):
    llm = ChatGoogleGenerativeAI(model=model, temperature=0)
    parser = PydanticOutputParser(pydantic_object=InsuredDetails)

    prompt = PromptTemplate(
        template="""
    Extract the insured business details from the following documents.
    Make sure to extract the full address. If the address is spread across multiple lines
    (e.g., "building number" and "street name" separately),
    combine them to form a full address in the format: street name, building number, city.
    The contact phone number format should be: 05xxxxxxxx (ten digits).
    
    {format_instructions}

    ### Examples:

    Text:
    כתובת מבוטח: 9, טבריה
    הערות לכתובת המבוטח: תמר 9

    Output:
    business_name: null
    business_address: תמר 9, טבריה
    contact_name: null
    contact_phone: null
    ~~~~~~~~~~~~~~~~~~~~

    Text:
    Acme Corp is based at 123 Main St, Springfield. The contact person is Jane Doe, 055-123-4567.

    Output:
    business_name: Acme Corp
    business_address: 123 Main St, Springfield
    contact_name: Jane Doe
    contact_phone: 0551234567
    ~~~~~~~~~~~~~~~~~~~~

    ### Now extract the details for the following document:

    Text:
    {input_text}
    """,
        input_variables=["input_text"],
        partial_variables={"format_instructions": parser.get_format_instructions()},
    )

    chain = prompt | llm | StrOutputParser()

    results = []
    for doc in docs:
        # Use the invoke method instead of run
        raw_output = chain.invoke({"input_text": doc})
        parsed = parser.parse(raw_output)
        results.append(parsed)

    for item in results:
        print(item)

    df_results = parse_to_df(results)

    return results, df_results