import pandas as pd
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from dotenv import load_dotenv
from hanaflow.config import LLMConfig


class LLMService:
    def __init__(self, file_path):
        self.file_path = file_path
        self.df = None
        self.null_values = None
        self.duplicate_values = None
        self.shape = None
        self.columns = None

        # LLM setup
        self.llm_config = LLMConfig()
        self.llm_config.validate()
        self.llm = ChatGoogleGenerativeAI(
            model=self.llm_config.model,
            google_api_key=self.llm_config.api_key
        )
        self.parser = StrOutputParser()

    # ------------------------------------------------------------------
    # Step 1 — Load data
    # ------------------------------------------------------------------

    def load_transform_data(self):
        self.df = pd.read_csv(self.file_path)
        return self.df

    # ------------------------------------------------------------------
    # Step 2 — Gather basic stats
    # ------------------------------------------------------------------

    def gather_info(self):
        if self.df is None:
            raise RuntimeError("Call load_transform_data() first.")

        self.null_values = self.df.isnull().sum()
        self.duplicate_values = self.df.duplicated().sum()
        self.shape = self.df.shape
        self.columns = list(self.df.columns)

        return {
            "null_values": self.null_values,
            "duplicate_values": self.duplicate_values,
            "shape": self.shape,
            "columns": self.columns
        }

    # ------------------------------------------------------------------
    # Step 3 — Generate a plain-English summary of the dataset
    # ------------------------------------------------------------------

    def generate_summary(self) -> str:
        """
        Uses the LLM to generate a plain-English summary of the uploaded
        dataset — covering what the data is about, key columns, data quality,
        and potential ML use cases.
        """
        if self.df is None:
            raise RuntimeError("Call load_transform_data() first.")

        # Build a rich context string from the dataframe
        sample_rows = self.df.head(5).to_string(index=False)
        dtypes_info = self.df.dtypes.to_string()
        null_info   = self.df.isnull().sum().to_string()
        stats_info  = self.df.describe(include="all").to_string()

        prompt_template = PromptTemplate(
            input_variables=[
                "filename", "shape_rows", "shape_cols", "columns",
                "dtypes", "nulls", "duplicates",
                "sample", "stats"
            ],
            template="""
You are a data analyst assistant. A user has uploaded a dataset into SAP HANA.
Analyze the following dataset information and provide a clear, structured summary.

--- DATASET INFO ---
File     : {filename}
Shape    : {shape_rows} rows x {shape_cols} columns
Columns  : {columns}

--- DATA TYPES ---
{dtypes}

--- NULL VALUES PER COLUMN ---
{nulls}

--- DUPLICATE ROWS ---
{duplicates}

--- SAMPLE ROWS (first 5) ---
{sample}

--- STATISTICAL SUMMARY ---
{stats}

Please provide:
1. What this dataset is about (2-3 sentences)
2. Key columns and what they represent
3. Data quality issues (nulls, duplicates, type mismatches)
4. Potential ML use cases for this dataset
5. One quick insight from the data

Keep the response clear and beginner-friendly.
"""
        )

        chain = prompt_template | self.llm | self.parser

        shape = self.shape or self.df.shape

        response = chain.invoke({
            "filename"   : self.file_path,
            "shape_rows" : shape[0],
            "shape_cols" : shape[1],
            "columns"    : self.columns or list(self.df.columns),
            "dtypes"     : dtypes_info,
            "nulls"      : null_info,
            "duplicates" : self.duplicate_values if self.duplicate_values is not None
                           else self.df.duplicated().sum(),
            "sample"     : sample_rows,
            "stats"      : stats_info,
        })

        return response

    # ------------------------------------------------------------------
    # Step 4 — Ask any question about the dataset
    # ------------------------------------------------------------------

    def ask(self, question: str) -> str:
        """
        Lets the user ask any natural language question about the dataset.
        The LLM answers using the actual data context.

        Examples
        --------
        service.ask("What is the average salary by city?")
        service.ask("Which column has the most null values?")
        service.ask("Is this dataset suitable for classification?")
        """
        if self.df is None:
            raise RuntimeError("Call load_transform_data() first.")

        # Give LLM enough context to answer accurately
        sample_rows  = self.df.head(10).to_string(index=False)
        dtypes_info  = self.df.dtypes.to_string()
        stats_info   = self.df.describe(include="all").to_string()
        null_info    = self.df.isnull().sum().to_string()

        prompt_template = PromptTemplate(
            input_variables=[
                "question", "filename", "shape_rows", "shape_cols",
                "columns", "dtypes", "nulls",
                "sample", "stats"
            ],
            template="""
You are a helpful data analyst assistant. A user has uploaded a dataset into SAP HANA
and is asking a question about it.

--- DATASET CONTEXT ---
File     : {filename}
Shape    : {shape_rows} rows x {shape_cols} columns
Columns  : {columns}

--- DATA TYPES ---
{dtypes}

--- NULL VALUES PER COLUMN ---
{nulls}

--- SAMPLE DATA (first 10 rows) ---
{sample}

--- STATISTICAL SUMMARY ---
{stats}

--- USER QUESTION ---
{question}

Answer the question clearly and concisely based on the dataset information above.
If the answer requires actual computation beyond what is shown, say so and suggest
how the user can find the answer using pandas.
"""
        )

        chain = prompt_template | self.llm | self.parser

        response = chain.invoke({
            "question"   : question,
            "filename"   : self.file_path,
            "shape_rows" : self.df.shape[0],
            "shape_cols" : self.df.shape[1],
            "columns"    : list(self.df.columns),
            "dtypes"     : dtypes_info,
            "nulls"      : null_info,
            "sample"     : sample_rows,
            "stats"      : stats_info,
        })

        return response
