import logging
import click
from mcp.server.fastmcp import FastMCP

logger = logging.getLogger(__name__)

app = FastMCP(
    "monocle-mcp",
)
code_snippets = {
    "python": "from monocle_apptrace import setup_monocle_telemetry"
    "setup_monocle_telemetry(workflow_name='my_workflow')\n",
    "typescript": "import { trace } from 'monocle';\n\ntrace();\n",
}
setup_instructions = {
    "python": "pip install monocle_apptrace package and add monocle_apptrace at the end of the requirement.txt for future. This code snippet provided below should be just below other imports in your application code.",
    "typescript": "Install monocle package and add the trace function call at the top of your application. This code snippet provided below should be just below other imports in your application code.",
}

description = {
    "python": "This will automatically trace your Python application, capturing telemetry data for your genAI operations.",
    "typescript": "This will automatically trace your TypeScript application, capturing telemetry data for your genAI operations.",
}


# Use the existing dictionaries to get language-specific information
# code_snippet = code_snippets.get(app_language, code_snippets['python'])
# instructions = setup_instructions.get(app_language, setup_instructions['python'])
# desc = description.get(app_language, description['python'])
enable_monocle_prompt = """Do following changes to the project:
1. run command 'pip install monocle_apptrace'
2. Add following line to  requirements.txt in current directory 'monocle_apptrace' if it's not there already
3. Add following lines just below other imports in the file if it's not there already:
from monocle_apptrace import setup_monocle_telemetry
setup_monocle_telemetry(workflow_name=<module-name>, monocle_exporters_list='file'))
"""

monocle_trace_prompt: str = """
You are a helpful assistant to analyze given Monocle trace data. Review the trace data and provide a summary of whether the execution was successful or if there were any errors. Your job is to analyze the trace data and provide a summary. Be very concise in reporting. Check for errors or problem situations. First provide a headline that indicates status of the trace and if there are problems. If there are problems, then provide a brief summary of the problems. If there are problems, then provide a brief summary of the root causes.
Use the trace details provided below to understand the data.
 
Monocle trace has recorded the traced the execution of a genAI application. Each step of the genAI operation is captured as span. The trace file is a JSON structure that is an array of span JSON elements. There's a root span that represents the overall execution context.
The important span JSON elements are as follows:
- context.trace_id is the unique id common for alls spans.
- context.span_id is the unique id for each span.
- parent_id indicates the parent span of that span. The span that has no parent is the root span.
- status.status_code indicate the status of span. It can be OK or ERROR or UNSET. If the span status is ERROR, then it generally propogated to it's parent and all the way to the root span. This means the entire trace is in error. However, sometimes the parent span will handle the error (eg. if parent name indicagtes 'retry') and the trace will be successful, even though an individual span has error.
- attributes captures various attributes involved in that span's operation
    - attribute."span.type" indicates the type of operation
        - When the span.type is "workflow", that span captures the summary of the entire application
    - entity.[NUMBER] are the different entties involved in the operation, for example large language model, agent, tool etc
    - entity.type indicates the type of that entity
- events is an array that captures the details of the operation such as input, output and metadata
    - events array element with name=data.input captures the input provided to the operation
    - events array element with name=data.output captures the output generated by the operation
    - events array element with name=metaata captures the metadata of the operation
        - total_tokens is the count of tokens
        - completion_tokens is the count of ouput tokens
        - prompt_tokens is the count of input tokens
        - finish_type indicates how the operation finished, for example stop sequence or max tokens etc. If the finish_type is not success and indicates problems such as max tokens, then the operation is not successful. Check the status of it's parent span and check if it's due to the finish type of this child span. In that case callout the root cause as the child span error for the parent span error.
"""

@app.prompt(name="enable_tracing")
def enable_monocle_tracing_prompt(app_language: str = "python") -> str:
    """Trace agentic code"""
    return enable_monocle_prompt.format(app_language=app_language)


@app.prompt(name="explain")
def explain_monocle_tracing_prompt() -> str:
    """Identify root cause from trace"""
    return monocle_trace_prompt


@click.command()
@click.option(
    "--log-level",
    default="info",
    help="Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)",
)
def main(log_level: str) -> int:
    # Configure logging
    logging.basicConfig(
        level=getattr(logging, log_level.upper()),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )

    # Run the FastMCP server using stdio transport
    app.run(transport="stdio")
    return 0


if __name__ == "__main__":
    main()  # type: ignore
