from pydantic import BaseModel

from wandern.agents.base_agent import AgentResponse
from wandern.graph import MigrationGraph
from wandern.models import Config

from .base_agent import BaseAgent


class MigrationSQL(BaseModel):
    up_sql: str | None
    down_sql: str | None
    message: str | None


class MigrationAgentResponse(AgentResponse[MigrationSQL, str]):
    pass


SYSTEM_PROMPT = """
You are a helpful assistant who helps users generate SQL queries based on their requests.

You are tasked with generating the necessary components for a SQL migration file as per the user prompt.
The output of the request will be of type `MigrationAgentResponse` which has 3 parts.
- `data` - Containing the data generated by the agent.
- `message` - A short message from the agent to the user.
- `error` - An error message, if applicable.

Below is the json structure of the output:
```json
{
    "data": {
        "up_sql": "string | null",
        "down_sql": "string | null",
        "message": "string | null"
    },
    "message": "string",
    "error": "string | null"
}
```

Below is the python data models (in Pydantic):
```python
class MigrationSQL:
    up_sql: str | None
    down_sql: str | None
    message: str | None

class MigrationAgentResponse:
    data: MigrationSQL
    message: str
    error: str | None = None
```

The fields are defined as follows:

1. `data` - Contains the sql scripts you generated
    1. `up_sql`: The SQL statements to apply the migration.
    2. `down_sql`: The SQL statements to revert the migration.
    3. `message`: A very brief (10-12 characters) description of the migration. Should fit in one line.
2. `message`: A short message from the agent to the user.
3. `error`: An error message, if applicable.

You are to generate ONLY the output of the request in the format of `MigrationAgentResponse`
and nothing more.
you need to provide a helpful message in the `message` property outlining the changes
you made to the user.
You will get the additional context of the SQL dialect in the ADDITIONAL CONTEXT.
We currently only support `sqlite` and `postgresql`.

IMPORTANT:
You must ensure that the SQL script generated is non-invasive and does not result in data loss
and should account for graceful migration of existing rows.
You must ensure that the generated SQL is valid for that dialect.
Format the SQL strings so that they are properly indented and newlined.
END OF IMPORTANT

You will get the additional context of all the previous revisions in descending order of
creation (most recent being first) in the ADDITIONAL CONTEXT.
"""


class MigrationAgent(BaseAgent[MigrationSQL, str]):
    def __init__(
        self,
        config: Config,
    ):
        super().__init__(
            role="migration assistant",
            task="generate SQL migration files",
            system_prompt=SYSTEM_PROMPT,
        )

        self.config = config
        self.graph = MigrationGraph.build(config.migration_dir)
        self.system_prompt = self.create_system_prompt(
            role="migration assistant",
            task="generate SQL migration files",
            system_prompt=SYSTEM_PROMPT,
        )

    @property
    def output_type(self):
        return MigrationSQL

    def generate_additional_context(self):
        additional_context = []

        revisions = list(self.graph.iter())
        revisions = revisions[::-1]  # sort by most recent

        for rev in revisions:
            additional_context.append(rev.model_dump_json(indent=2))

        additional_context_str = "\n".join(additional_context)
        additional_context_str = (
            f"PAST REVISIONS:\n{additional_context_str}\n"
            if additional_context_str
            else ""
        )
        additional_context_str += f"SQL DIALECT: {self.config.dialect}\n"

        return additional_context_str

    def generate_revision(
        self,
        user_prompt: str,
    ):
        custom_prompt = self.create_structured_prompt(
            user_prompt=user_prompt,
            additional_context=self.generate_additional_context(),
        )

        response = super().run(
            custom_prompt,
        )

        return response
