import textwrap
from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import JSONResponse
from scribe_python_client.client import ScribeClient
from dotenv import load_dotenv
from pydantic import BaseModel
from typing import *
from pathlib import Path
from contextlib import asynccontextmanager
import logging
import json
import os


load_dotenv()
API_TOKEN = os.getenv("SCRIBE_TOKEN")
client = ScribeClient(api_token=API_TOKEN)

#test
toolDocsPath = Path(__file__).parent / "toolsDescriptions/tool_docs.json"

class Filter(BaseModel):
    col: str
    op: str
    val: Any

class Metric(BaseModel):
    label: str
    expressionType: str
    sqlExpression: Optional[str] = None

class VulnerabilityQuery(BaseModel):
    column: Optional[List[str]] = []
    filters: Optional[List[Filter]] = []
    metrics: Optional[List[Metric]] = []
    orderby: Optional[List[List[Any]]] = []
    row_limit: Optional[int] = 20

class QueryResponse(BaseModel):
    format: str
    status: bool
    status_info: str
    json_data: object
    str: str


def get_products_list() -> list:
    result = client.get_products()
    listOfProducts = []
    for i in range(len(result)):
        listOfProducts.append(result[i]["name"])
    return listOfProducts

app = FastAPI()

@app.post("/generate_query", operation_id="generateQuery")
def build_llm_prompt_from_docs( input: str = Query(...), function: str = Query(...)) -> JSONResponse:
    """
    Given tool documentation and a user question, generate a prompt
    for the LLM to decide which tool to call and how.
    Function Should be a valid query function name
    """

    """
    Refer to tool_docs.json for full schema details. Create a valid Input for a query based off the given string and then directly make the query.
    """
    with open(toolDocsPath, "r") as f:
        doc_section = json.load(f)
    tools_str = json.dumps(doc_section, indent=2)
    

    prompt = textwrap.dedent(f"""
        You are a tool selector and payload generator.
        Given the following tools and their documentation, choose the most appropriate tool to answer the user question. Then generate a valid JSON payload to call that tool.

        TOOLS:
        {tools_str[function]}

        USER QUESTION:
        {input}

        Respond ONLY in this JSON format:
        {{
            "tool_name": "<tool_name>",
            "payload": {{
                "columns": [...],
                "filters": [...],
                "metrics": [...],
                "groupby": [...],
                "orderby": [...],
                "row_limit": <int>
            }}
        }}
    """)
    
    return JSONResponse(content={"prompt": prompt.strip()})
        
@app.get("/products")
async def get_products():
    """
    Endpoint: Retrieve a list of products managed in Scribe: Should be called when asked for scribe products or a list of products.

    Args: Nothing

    Returns:
        JSONResponse: Json list of products
    """
    try:
        return client.get_products()
    except Exception as e:
        logging.exception(f"Error in get_products endpoint {e}")

@app.get("/datasets")
async def get_datasets():
    """
    Endpoint: Retrieve all datasets: This should be called when anything about datasets is mentions.

    Args: Nothing

    Returns:
        JSONResponse: Json list of datasets
    """
    try:
        return client.get_datasets()
    except Exception as e:
        logging.exception(f"Error in get_datasets endpoint: {e}")

@app.get("/{product_name}/vulns")
async def get_product_vulnerability(product_name: str):
    """
    Endpoint: Retrieve vulnerabilities for a specific product: This should be called when a specifc product and the user is asking for vulnerabilities.

    Args: product name.

    Returns:
        JSONResponse: Json list of product vulnerabilities
    """
    products = get_products_list()
    if product_name not in products:
        logging.exception(f"Product name: {product_name} is not a valid product")
        raise HTTPException(status_code=404, detail=str("Product not found"))
    try:
        return client.get_product_vulnerabilities(product_name)
    except Exception as e:
        logging.exception(f"Error in get_product_vulnerabilities endpoint: {e}")
    

@app.get("/{product_name}/policy")
async def get_product_policy(product_name: str):
    """
    Endpoint: Retrieve policy for a specific product: This should be called when a specifc product and the user is asking for policy.

    Args: product name.

    Returns: JSONResponse: Json list of product policy
    """
    products = get_products_list()
    if product_name not in products:
        logging.exception(f"Product name: {product_name} is not a valid product")
        raise HTTPException(status_code=404, detail=str("Product not found"))
    try:
        return client.get_policy_results(product_name)
    except Exception as e:
        logging.exception(f"Error in get_policy_results endpoint: {e}")

@app.post("/vulnerabilities")
async def query_vulnerabilities(input: dict):
    """
    This tool is used to answer questions about vulnerabilities like: What are the vulnerabilities of my product? What are the vulnerabilities of a specific component? What are the vulnerabilities of a specific version of a component? Which products have a specific vulnerability, or a vulnerability with a specificy characteristic (such as has_key)

    Args: 
    {
    "column": [],
    "filters" [],
    "metrics" [],
    "orderby": [],
    "row_limit": 20
    }
    - "columns": list of columns to select. Use only these valid columns:
    [('component_id', 'Dependency ID'), ('vulnerability_id', 'ID'), ('vector', 'Vector'), ('severity_string', 'Severity String'), ('base_severity_string', 'Base Severity String'), ('advisory_status', 'Advisory Status'), ('final_severity', 'Final Severity'), ('tool', 'Tool'), ('teamId', 'Teamid'), ('pipeline_run_id', 'Pipeline Run Id'), ('pkg', 'Pkg'), ('latest_package_version', 'Dependency Latest Version'), ('attestation_ids', 'Attestation IDs'), ('cvss_score', 'CVSS'), ('logical_app', 'Product'), ('vul_component_fixed_in_versions', 'Vulnerability Fix Version'), ('will_not_fix', 'Will Not Fix'), ('attestation_name', 'Attestation Name'), ('tag', 'Component Version'), ('targetName', 'Component'), ('base_score', 'Base Score'), ('source_layer', 'Layer'), ('text', 'Text'), ('is_latest_logical_version', 'Latest product version'), ('update', 'Update'), ('cwe_array', 'CWEs'), ('has_fix', 'Has Fix'), ('component_purl', 'Dependency PURL'), ('component_version', 'Dependency Version'), ('component_locations', 'Dependency Locations'), ('vul_date_changed', 'Vulnerability Last Change'), ('vul_component_created', 'Dependency Date'), ('epssProbability', 'EPSS'), ('vendor', 'Vendor'), ('direct_cwes', 'Direct Cwes'), ('advisory_modified', 'Advisory Modified'), ('hyper_links', 'Hyper Links'), ('advisory_text', 'Advisory Text'), ('cwes', 'CWE_Object'), ('advisory_justification', 'Advisory Justification'), ('component_dependency_type', 'Component Dependency Type'), ('cpes', 'Dependency CPEs'), ('has_kev', 'Has KEV'), ('version_timestamp', 'Version Timestamp'), ('severity', 'Original Severity'), ('epssPercentile', 'Epsspercentile'), ('epssScoreDate', 'Epssscoredate'), ('vul_adv_id', 'Vul Adv Id'), ('dependency_analysis', 'Dependency Analysis'), ('main_tool', 'Main Tool'), ('vul_published_on', 'Published'), ('vul_hyperlinks', 'References'), ('component_name', 'Dependency'), ('product_id', 'Product Id'), ('labels', 'Labels'), ('vul_source', 'Source'), ('targets', 'Targets'), ('logical_app_version', 'Version'), ('id', 'Id')]

    - "filters": list of filter objects with keys:
        - "col": column name (must be valid)
        - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null
        - "val": filter value

    - "metrics": list of metrics with keys:
        - "label": output column name
        - "expressionType": "SQL" or "SIMPLE"
        - "sqlExpression": e.g. "COUNT(*)"

        - "groupby": list of valid column names

        - "orderby": list of [column_name, is_ascending_boolean]

         "row_limit": integer, max 30

        Output ONLY the JSON object with no extra text.
    """
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        result_str = client.query_vulnerabilities(json.dumps(input))
        logging.debug(result_str)
        if result_str.startswith("Error"):
            logging.exception("Error in query_vulnerabilities endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_vulnerabilities endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response
 
@app.post("/products")
async def query_products(input: dict):
    """
    To answer questions about products like: what are the/my products? What are the components of my product (the this case use the targetName column!)? What are the depdencieis of my product (search the component_name column)? Note that repositories and namespaces are part of the lineage table.

    Args: 
    {
    "column": [],
    "filters" [],
    "metrics" [],
    "orderby": [],
    "row_limit": 20
    }
    - "columns": list of columns to select. Use only these valid columns:
    [('more', ' '), ('license_is_permissive_na', 'License Is Permissive Na'), ('teamId', 'Teamid'), ('component_version', 'Dep version'), ('license_expression', 'Licesne expression'), ('pkg', 'Package Manager'), ('rstatus', 'Version Status'), ('component_purl', 'PURL'), ('dependency_parents', 'Dependency Parents'), ('dependency_children', 'Dependency Children'), ('component_dependency_type', 'Component Dependency Type'), ('logical_app', 'Product'), ('logical_app_version', 'Version'), ('scorecard_id', 'Scorecard Id'), ('version_is_up_to_date', 'Up-to-date'), ('product_id', 'Product Id'), ('attestation_name', 'Attestation Name'), ('high_severity_cves', 'C+H Vulns'), ('maintenance_score', 'Maintenance Score'), ('component_name', 'Dependency'), ('labels', 'Labels'), ('licenses', 'Licenses'), ('attestation_ids', 'Attestation Ids'), ('component_external_ref', 'Component External Ref'), ('license_expression_na', 'License Expression Na'), ('component_publisher', 'Component Publisher'), ('component_description', 'Component Description'), ('targets_translated', 'SBOM Type Translated'), ('component_author', 'Component Author'), ('relation', 'Relation'), ('main_tool', 'Main Tool'), ('maintenance_score_int', 'Maintenance Score'), ('block', 'Dependency'), ('tools', 'Tools'), ('license_is_permissive', 'Is permissive'), ('scorecard_score', 'OpenSSF score'), ('targets', 'SBOM Type'), ('component_type', 'Component type'), ('source_layer', 'Layer'), ('component_locations', 'Component Locations'), ('component_id', 'Component Id'), ('id', 'Id'), ('base_layer', 'Base Image'), ('latest_version', 'Dep latest ver'), ('tag', 'Component version'), ('targetName', 'Component'), ('component_group', 'Component Group'), ('version_timestamp', 'Version Timestamp'), ('attestation_ids_cnt', 'Attestation Ids Cnt'), ('pipeline_run_id', 'Pipeline Run Id')]
    - "filters": list of filter objects with keys:
        - "col": column name (must be valid)
        - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null
        - "val": filter value

    - "metrics": list of metrics with keys:
        - "label": output column name
        - "expressionType": "SQL" or "SIMPLE"
        - "sqlExpression": e.g. "COUNT(*)"

        - "groupby": list of valid column names

        - "orderby": list of [column_name, is_ascending_boolean]

         "row_limit": integer, max 30

        Output ONLY the JSON object with no extra text.
    """
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        result_str = client.query_products(json.dumps(input))
        if result_str.startswith("Error"):
            logging.exception("Error in query_products endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_products endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response

@app.post("/policy")
async def query_policy_results(input: dict):
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        result_str = client.query_policy_results(json.dumps(input))
        if result_str.startswith("Error"):
            logging.exception("Error in query_policy_results endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_policy_results endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response
    
@app.post("/lineage")
async def query_lineage(input: dict):
    """
    To answer software-factory related questions such as: What repos (or code repositories) are related/part of a product? or What are all the software-factory related components of a product (in this case query for asset_type and asset_name for the logical_app and logical_app_version)?
    Args: 
    {
    "column": [],
    "filters" [],
    "metrics" [],
    "orderby": [],
    "row_limit": 20
    }
    - "columns": list of columns to select. Use only these valid columns:
    [('platform_type', 'Platform Type'), ('properties', 'Info'), ('version_id', 'Version Id'), ('product_id', 'Product ID'), ('parent_name', 'Parent Name'), ('parent_type', 'Parent Type'), ('asset_name', 'Asset'), ('path', 'Path'), ('parent_id', 'Parent Id'), ('child_id', 'Child Id'), ('teamId', 'Teamid'), ('owner', 'Owner'), ('platform_name', 'Platform'), ('asset_type', 'Type'), ('external_id', 'External ID'), ('timestamp', 'Updated'), ('uri', 'URI'), ('prent_bloack', 'Parent'), ('more', ' More'), ('tooltip', 'Tooltip'), ('parent_external_id', 'Parent External ID'), ('logical_app', 'Product'), ('logical_app_version', 'Version'), ('asset_block', 'Asset'), ('g-source', 'G-Source'), ('g-target', 'G-Target')]    - "filters": list of filter objects with keys:
        - "col": column name (must be valid)
        - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null
        - "val": filter value

    - "metrics": list of metrics with keys:
        - "label": output column name
        - "expressionType": "SQL" or "SIMPLE"
        - "sqlExpression": e.g. "COUNT(*)"

        - "groupby": list of valid column names

        - "orderby": list of [column_name, is_ascending_boolean]

         "row_limit": integer, max 30

        Output ONLY the JSON object with no extra text.
    """
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        print(input)
        result_str = client.query_lineage(json.dumps(input))
        if result_str.startswith("Error"):
            logging.exception("Error in query_lineage endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_lineage endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response

@app.post("/risk")
async def query_risk(input: dict):
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        print(input)
        result_str = client.query_risk(json.dumps(input))
        if result_str.startswith("Error"):
            logging.exception("Error in query_risk endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_risk endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response

@app.post("/findings")
async def query_findings(input: dict):
    response = QueryResponse(
        format="",
        status=True,
        status_info="Query executed successfully",
        json_data={},
        str=""
    )
    try:
        print(input)
        result_str = client.query_findings(json.dumps(input))
        if result_str.startswith("Error"):
            logging.exception("Error in query_findings endpoint")
            raise HTTPException(status_code=400, detail=result_str)
        if result_str == f"\nNo data found":
            response.status = False
            response.status_info = "No data found"
        elif format == "json":
            response.json_data = result_str
        else:
            response.format = "str"
            response.str = result_str
    except HTTPException as he:
        raise he    
    except Exception as e:
        logging.exception(f"Error in query_findings endpoint: {e}")
        response.status = False
        response.status_info = str(e)
    return response